diff --git a/README.md b/README.md index 2926d51308bbfee3a3e947d85c03116fe58382e8..d3c7f5708c0e419b8ca50f2ea588fd8a3aac7855 100644 --- a/README.md +++ b/README.md @@ -26,18 +26,36 @@ geometry, Qwen material, completion, and PBR asset export stages. Target runtime: Python 3.10, PyTorch 2.8.0 + CUDA 12.8, torchvision 0.23.0, and SDPA attention. The Space compiles pinned public revisions of CuMesh, -FlexGEMM, nvdiffrast, and O-Voxel for NVIDIA L4 (sm_89), and installs the +FlexGEMM, nvdiffrast, and O-Voxel for the ZeroGPU CUDA architecture, and installs the matching official NATTEN 0.21.1 wheel. CUDA compiler archives are downloaded remotely from NVIDIA and checked against their SHA256 digests. -The live Space uses one dedicated NVIDIA L4 (24 GB). GPU usage is billed to -the Space owner; visitors do not need ZeroGPU quota. The retained `spaces.GPU` -decorators are inactive on dedicated hardware. L4 wheels use a separate -`/data/native-wheels/torch280-cu128-cp310-sm89-v1` cache. - -The ZeroGPU source is preserved at tag -`zerogpu-before-l4-20261020-77f93e49`, along with its original sm_120 cache. -Rollback requires restoring that source and choosing ZeroGPU hardware. +The live Space uses ZeroGPU. Visitors sign in with their own Hugging Face +accounts and submit GPU operations through the official Gradio client in the +Hugging Face Space page. The client's signed quota header stays attached to the +live Gradio request; the Space owner's model-download token is never supplied +as a visitor's GPU credential. Missing quota headers and explicit quota-token +errors are rejected before inference. + +Opening the page and starting the Space do not load models or compile native +extensions. An explicitly submitted task starts its required CPU preparation. +Background removal loads only the pinned RMBG model, independently of native +compilation and the full 18-model 3D pipeline. Its upload panel shows one status +message throughout image loading and background removal. 3D preparation shows +per-phase percentages from downloaded bytes, loaded models, or completed build steps. +Foreground removal, generation, export, +and preview recovery each acquire a temporary `spaces.GPU` lease. No permanent +CUDA worker, unattended GPU queue, curated refresh, or GPU keepalive runs in +ZeroGPU mode. Saved unfinished requests require an explicit Resume action; +opening history never starts GPU work. Independent requests can execute +concurrently, subject to Hugging Face GPU availability. Repeated submissions of +the same operation execute once. +Native wheels use a separate ZeroGPU architecture cache under +`/data/native-wheels`; the former L4 cache is retained for dedicated hardware. + +The application mounts Gradio into FastAPI instead of calling `Blocks.launch`. +With the pinned `spaces==0.51.3`, it explicitly reports SDK startup after GPU +functions are registered. CPU-only preparation starts later, on demand. An isolated L4 validation on 2026-10-20 completed seed 42 pitcher-plant generation in 338.2 seconds and GLB export in 49.8 seconds. Sampled peak @@ -75,9 +93,13 @@ originals in the persistent Bucket are retained. ## Startup and checkpoint selection Mount a writable Bucket at `/data` and provide an approved read token as the -private `HF_TOKEN` Space secret. Native extension installation and a shared -background model-preparation worker start automatically with the server. -The worker waits for native runtime readiness before loading weights. +private `HF_TOKEN` Space secret. In ZeroGPU mode, only the HTTP interface and +history storage initialize with the server. A verified task's explicit prepare +request starts native extension installation and one shared CPU model-loading +controller for 3D tasks. Background-removal tasks use their own one-model CPU +controller and do not start native installation or full-pipeline loading. The +3D controller waits for native runtime readiness before loading weights. +Preparation does not acquire a GPU lease. Successful CUDA validation reports are bound to the wheel inventory and cached in the Bucket. Matching cached wheels and checks are reused after restarts. Diagnostic panels and GPU test buttons are not part of the public @@ -92,7 +114,9 @@ during startup. Archived project weights remain under shared across all user sessions and generations. A request arriving during startup waits for that same worker, with weight-loading progress in the loading panel. Opening or refreshing a page does not start another -load. A restart or wake from sleep automatically starts preparation again. +load. After a restart or wake from sleep, CPU preparation waits for the next +explicit task. Browsing history, selecting an existing preview, and status polls +never initiate model loading. Transfers retry without discarding their verified source identity or partial bytes; exhausted retries report failure instead of an endless loading state. Private artifact reads have their own limited worker capacity so slow storage @@ -127,6 +151,12 @@ Optional Space variables: - `SUPRATEX_MODEL_REPO`: project checkpoint repository, default `supra3d/supra-tex`. - `SUPRATEX_MODEL_REVISION`: immutable project checkpoint commit SHA. +- `SUPRATEX_ZERO_GENERATE_SECONDS`, `SUPRATEX_ZERO_FOREGROUND_SECONDS`, + `SUPRATEX_ZERO_EXPORT_SECONDS`, `SUPRATEX_ZERO_RECOVER_SECONDS`: requested GPU + durations, defaulting to 180, 60, 180, and 120 seconds respectively. The ZeroGPU + scheduler may apply a hardware duration factor. These are time limits, not + measured completion times; insufficient personal quota rejects a request. + The frontend waiting estimate is set to 390 seconds in `app.py`. This estimate is not measured model progress or an inference time limit. diff --git a/app.py b/app.py index da46d61a8e857e57d1c2ae130bc14f42c98a4269..8e02d98401682cd4a0cffd101a234144ca09ec97 100644 --- a/app.py +++ b/app.py @@ -1,9 +1,21 @@ -"""Supra-Tex with account-owned history and browser-independent queued tasks.""" +"""Supra-Tex with private history and on-demand per-user GPU operations.""" import os from pathlib import Path import sys import threading +# Install ZeroGPU's CUDA interception before any project/Torch imports. +from hf_space.gpu_mode import enabled as zero_gpu_enabled +if zero_gpu_enabled(): + # The SDK sets PyTorch's one-shot inter-op count during its first import. + os.environ.setdefault('OMP_NUM_THREADS', '1') + os.environ.setdefault('MKL_NUM_THREADS', '1') +import spaces +if zero_gpu_enabled(): + # HF checks registered GPU tasks as soon as the HTTP server starts. Model + # weights remain unloaded until a visitor submits a task. + from hf_space import zerogpu_runtime + from scripts.run_example import ROOT as SOURCE_ROOT, _pipeline_config sys.path.insert(0, str(_pipeline_config().runtime_root(SOURCE_ROOT)/'utils3d')) @@ -17,7 +29,8 @@ from hf_space.history_routes import install_routes from hf_space.history_ui import build_ui, CSS, THEME, I18N ROOT = Path(__file__).resolve().parent -PUBLIC_PATHS = [ROOT/'hf_space/gallery-viewer.bundle.js', ROOT/'hf_space/gallery-lights'] +PUBLIC_PATHS = [ROOT/'hf_space/gallery-viewer.bundle.js', ROOT/'hf_space/gallery-lights', + ROOT/'hf_space/vendor/gradio-client-2.1.0.js'] PRIVATE_PATHS = ['/data', '/tmp/supratex-jobs', '/tmp/gradio', '/tmp/supratex-history-private', '/tmp/supratex-progress'] @@ -25,8 +38,16 @@ PRIVATE_PATHS = ['/data', '/tmp/supratex-jobs', '/tmp/gradio', def create_app(service=None): service = service if service is not None else Service() demo = build_ui(deployment_status) - # Browser callbacks only mount the UI. All GPU operations share Service.runner. - demo.queue(api_open=False) + from hf_space.gpu_mode import enabled + if enabled(): + from hf_space.zerogpu_jobs import register + register(demo, service) + # ZeroGPU callbacks retain each visitor's request through GPU scheduling. + demo.queue(api_open=False, max_size=None) + if enabled(): + # Gradio 6.10 changes None to 1 on ZeroGPU. Restore its Queue's actual + # unlimited value; event transport must not add a global admission cap. + demo._queue.max_size = None web = gr.mount_gradio_app(FastAPI(), demo, path='/', allowed_paths=[str(path) for path in PUBLIC_PATHS], blocked_paths=PRIVATE_PATHS, css=CSS, theme=THEME, i18n=I18N, ssr_mode=False, footer_links=['api', 'gradio']) @@ -36,17 +57,20 @@ def create_app(service=None): install_routes(gradio_app, service.store, deployment_status=deployment_status, examples_provider=service.examples, example_file_provider=service.example_file, prepare_example=service.prepare_example, preview_provider=service.preview, + prepare_models=service.prepare_models, public_gradio_paths=PUBLIC_PATHS, private_file_cache=service.private_files) web.state.history_service = service return web def main(): - # Do not block the HTTP server on Bucket reads or weight/model preparation. - # Public models normally arrive directly from their immutable CDN URLs. - # Prepare a local fallback only when explicitly requested, never at boot. - native_build.start() - release_backend.start_preparation() + if zero_gpu_enabled(): + zerogpu_runtime.start_mounted_app() + else: + # Dedicated GPU deployments retain their startup preparation. ZeroGPU + # starts this shared CPU preparation only when a visitor submits work. + native_build.start() + release_backend.start_preparation() service = Service() web = create_app(service) service.start() diff --git a/deployment/vendor/eigen-21e4582d1739107337a03460c81412981130373e.tar.gz b/deployment/vendor/eigen-21e4582d1739107337a03460c81412981130373e.tar.gz new file mode 100644 index 0000000000000000000000000000000000000000..326824fc150151ae1dc106934329bc5206535abf --- /dev/null +++ b/deployment/vendor/eigen-21e4582d1739107337a03460c81412981130373e.tar.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:399bbf2c2239b286f21edfebfe1430a689c0995acd32e8017701346f6d35be7a +size 3025643 diff --git a/deployment/vendor/eigen-source.json b/deployment/vendor/eigen-source.json new file mode 100644 index 0000000000000000000000000000000000000000..c014d5ebc128e49a42862911015ecd9e1f8e8c18 --- /dev/null +++ b/deployment/vendor/eigen-source.json @@ -0,0 +1,18 @@ +{ + "project": "Eigen", + "repository": "https://gitlab.com/libeigen/eigen.git", + "revision": "21e4582d1739107337a03460c81412981130373e", + "source": "https://gitlab.com/libeigen/eigen/-/archive/21e4582d1739107337a03460c81412981130373e/eigen-21e4582d1739107337a03460c81412981130373e.tar.gz", + "archive": "eigen-21e4582d1739107337a03460c81412981130373e.tar.gz", + "sha256": "399bbf2c2239b286f21edfebfe1430a689c0995acd32e8017701346f6d35be7a", + "bytes": 3025643, + "required_by": { + "repository": "https://github.com/microsoft/TRELLIS.2.git", + "revision": "5565d240c4a494caaf9ece7a554542b76ffa36d3", + "submodule": "o-voxel/third_party/eigen" + }, + "license": "Primarily MPL-2.0; some files have BSD or other MPL-2.0-compatible licenses, as documented by upstream COPYING.README.", + "license_files": ["LICENSE", "COPYING.README", "COPYING.MPL2", "COPYING.BSD", "COPYING.APACHE", "COPYING.MINPACK"], + "modified": false, + "note": "Unmodified official source archive. All upstream copyright notices and licenses remain in the archive and extracted source. Stored with Hugging Face LFS and included in the Space checkout to avoid a build-time GitLab fetch." +} diff --git a/hf_space/backend.py b/hf_space/backend.py index ea0cdbc3a86ca9192aa7ef05f47c3cbde1ff49d3..548ffaf288938793aaec028f2864f6a8611ef337 100644 --- a/hf_space/backend.py +++ b/hf_space/backend.py @@ -219,7 +219,7 @@ def generate(image, seed, progress=None, timeout_seconds=None, generation_token= job = Path('/tmp/supratex-jobs') / uuid.uuid4().hex job.mkdir(parents=True) generation_progress.update(generation_token, status='running', job=job.name, - label='Initializing the pipeline', generation_started_unix=time.time()) + label='Starting generation…', generation_started_unix=time.time()) image.save(job / 'input.png') config = str(READY / 'models.json') output = job / 'asset.npz' diff --git a/hf_space/brand/huggingface-logo.svg b/hf_space/brand/huggingface-logo.svg new file mode 100644 index 0000000000000000000000000000000000000000..c43066896c2d5c7d56d40d1576c3ffb5553589b9 --- /dev/null +++ b/hf_space/brand/huggingface-logo.svg @@ -0,0 +1,37 @@ + + + + + + + + + + + diff --git a/hf_space/curated_refresh.py b/hf_space/curated_refresh.py index 0549cf66e6630778bde6692c9b676bdd8a08fbdc..1c635109ff3c573917969d1183e4199b883b36c8 100644 --- a/hf_space/curated_refresh.py +++ b/hf_space/curated_refresh.py @@ -169,7 +169,7 @@ class Coordinator: raise RefreshError('The approved input pixels changed.') upload = self.store.save_upload(self.owner, source.read_bytes()) job = self.store.submit(self.owner, upload['id'], self.plan['seed'], - self._key(entry['name'], 'generate')) + self._key(entry['name'], 'generate'), workflow=None) self.state['items'][entry['name']] = { 'name': entry['name'], 'source': deepcopy(entry), 'job_id': job['id'], 'upload_id': job['upload_id'], 'input_sha256': job['input_sha256'], diff --git a/hf_space/generation_progress.py b/hf_space/generation_progress.py index ed075c7660114537d13f00bf33a17b261bb744af..88016ce609c21439392d6c26c48b09a267f48b91 100644 --- a/hf_space/generation_progress.py +++ b/hf_space/generation_progress.py @@ -54,13 +54,19 @@ def start(*, preparing=False): def stage_label(name): if name in ('pipeline', 'frontend'): - return 'Initializing the pipeline' + return 'Starting generation…' if name in ('frontend.validation', 'frontend.preprocess'): return 'Preparing the image' if name == 'frontend.camera': return 'Estimating the camera' if name == 'geometry.model_load': return 'Loading the geometry model' + if name == 'geometry.save': + return 'Saving geometry checkpoint' + if name == 'geometry.preview': + return 'Rendering clay preview' + if name == 'geometry.cleanup': + return 'Cleaning the geometry' if name == 'frontend.geometry' or name.startswith('geometry.'): return 'Generating geometry' if name in ('qwen.checkpoint_load', 'qwen.transformer_load', 'qwen.adapter_load'): @@ -74,6 +80,8 @@ def stage_label(name): return 'Loading completion models' if name in ('sample', 'sample_25_steps'): return 'Completing surface materials' + if name == 'texture.completion': + return 'Completing surface materials' if name in ('decoder_load', 'decode'): return 'Decoding the 3D asset' if name == 'ovoxel_save': diff --git a/hf_space/geometry_preview.py b/hf_space/geometry_preview.py new file mode 100644 index 0000000000000000000000000000000000000000..866d34dcb7cc5e56ad3c3e9fd6dfb338f7fe4daa --- /dev/null +++ b/hf_space/geometry_preview.py @@ -0,0 +1,52 @@ +"""Lightweight clay views of a saved mesh, without material inference or UVs.""" +import json +from pathlib import Path + +from hf_space.timing import span + + +def render_geometry(folder, vertices, faces, camera): + import numpy as np + import torch + from PIL import Image + from supratex.representations import Mesh + from supratex.renderers import MeshRenderer + from supratex.utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics + from hf_space.ovoxel_preview import INPUT_VIEW_LAYOUT, orbit_parameters + + folder = Path(folder) + with span('geometry.preview'): + preview = Mesh(vertices.cuda(), faces.cuda()) + if len(faces) > 16777216: + preview.simplify(16777216) + count, resolution = 16, 768 + render_camera = {'camera_angle_x': float(camera['camera_angle_x']), + 'distance': float(camera['distance']) * float(camera.get('mesh_scale', 1.0))} + yaws, pitches, distance, fov = orbit_parameters(count, render_camera) + extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitches, distance, fov) + rotation = torch.tensor([[1.,0,0,0],[0,0,-1.,0],[0,1.,0,0],[0,0,0,1]], device='cuda') + renderer = MeshRenderer({'resolution': resolution, 'ssaa': 1, + 'near': min(.01, distance/100), 'far': max(10., distance+2)}) + views = folder / 'views' + views.mkdir(exist_ok=True) + light = np.asarray([-.35, .45, 1.], dtype=np.float32) + light /= np.linalg.norm(light) + with torch.no_grad(): + for index, (extrinsic, intrinsic) in enumerate(zip(extrinsics, intrinsics)): + result = renderer.render(preview, extrinsic @ rotation, intrinsic, + return_types=['normal', 'mask']) + normal = result['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 2 - 1 + mask = result['mask'].detach().cpu().numpy() + if mask.ndim == 3: + mask = mask[0] + # A neutral, view-relative diffuse light reveals shape only. + diffuse = np.clip(np.sum(normal * light, axis=-1), 0, 1) + gray = (70 + 150 * diffuse) * mask + 25 * (1 - mask) + pixels = np.repeat(np.clip(gray, 0, 255).astype(np.uint8)[..., None], 3, axis=-1) + Image.fromarray(pixels).save(views / f'clay-{index:02d}.webp', quality=90) + (folder / 'preview.json').write_text(json.dumps({ + 'views': count, 'modes': ['clay'], 'default_mode': 'clay', 'resolution': resolution, + 'device': 'cuda', 'renderer': 'MeshRenderer', 'lighting': 'Neutral diffuse clay', + 'source_faces': len(faces), 'preview_faces': len(preview.faces), + 'view_layout': INPUT_VIEW_LAYOUT, 'input_view': 0, 'input_camera': render_camera, + 'automatic_glb_export': False, 'output_kind': 'geometry'}), encoding='utf-8') diff --git a/hf_space/geometry_results.py b/hf_space/geometry_results.py new file mode 100644 index 0000000000000000000000000000000000000000..c6285fd20569220689e6ea0e1b9f66b323b13d64 --- /dev/null +++ b/hf_space/geometry_results.py @@ -0,0 +1,105 @@ +"""Durable, integrity-checked geometry handoffs for explicit texture requests.""" +import hashlib +import json +from pathlib import Path +import shutil + +from hf_space import ovoxel_results +from hf_space.split_pipeline import CHECKPOINT_FILES, OPTIONAL_CHECKPOINT_FILES, validate_checkpoint + +HANDOFF_FILES = ('input.png', 'geometry.npz', + *(f'geometry-checkpoint/{name}' for name in CHECKPOINT_FILES)) + + +def _handoff_files(folder): + metadata = validate_checkpoint(folder / 'geometry-checkpoint') + return (*HANDOFF_FILES, *(f'geometry-checkpoint/{name}' for name in OPTIONAL_CHECKPOINT_FILES + if name in metadata['files'])) + + +def _digest(path): + digest = hashlib.sha256() + with path.open('rb') as stream: + for block in iter(lambda: stream.read(1024 * 1024), b''): + digest.update(block) + return digest.hexdigest() + + +def _file(root, name): + path = root / name + if (root.is_symlink() or path.is_symlink() or path.parent.is_symlink() + or not path.resolve().is_relative_to(root.resolve()) or not path.is_file() + or path.stat().st_size <= 0): + raise ValueError('The saved geometry is incomplete or invalid.') + return path + + +def persist(source, report): + """Publish the white preview and its exact texturing input; ready is last.""" + source = Path(source) + folder = source.parent + identifier = folder.name + destination = ovoxel_results.directory(identifier) + if source.name != 'geometry.npz' or destination.exists(): + raise ValueError('The geometry output destination is invalid.') + files, metadata = ovoxel_results.local_manifest(folder) + expected_frames = ovoxel_results._frame_names(metadata) + if metadata.get('modes') != ['clay'] or metadata.get('default_mode') != 'clay': + raise ValueError('Geometry must have a clay preview.') + if set(files) != {f'frames/{name}' for name in expected_frames}: + raise ValueError('The geometry preview is incomplete.') + # No user-selected path or archive member is accepted as a handoff file. + handoff_files = _handoff_files(folder) + names = (*handoff_files, 'report.json', 'preview.json', 'timings.jsonl', + *(f'views/{name}' for name in expected_frames)) + for name in names: + _file(folder, name) + destination.mkdir(parents=True) + bindings = {} + for name in names: + original = _file(folder, name) + target = destination / name + target.parent.mkdir(parents=True, exist_ok=True) + expected = {'size': original.stat().st_size, 'sha256': _digest(original)} + shutil.copyfile(original, target) + if target.stat().st_size != expected['size'] or _digest(target) != expected['sha256']: + raise ValueError('The saved geometry failed integrity validation.') + if name in handoff_files: + bindings[name] = expected + (destination / 'geometry-manifest.json').write_text(json.dumps({ + 'schema_version': 1, 'files': bindings}), encoding='utf-8') + (destination / 'display-manifest.json').write_text(json.dumps({ + 'files': files, 'preview_metadata': metadata}), encoding='utf-8') + (destination / 'ready.json').write_text(json.dumps({ + 'complete': True, 'identifier': identifier, 'output_kind': 'geometry', + 'project_revision': report['project_revision']}), encoding='utf-8') + return identifier + + +def restore(identifier, destination): + """Restore only a verified geometry handoff, without allocating a GPU.""" + folder = ovoxel_results.directory(identifier) + ready = json.loads(_file(folder, 'ready.json').read_text(encoding='utf-8')) + if ready.get('complete') is not True or ready.get('output_kind') != 'geometry' or ready.get('identifier') != identifier: + raise ValueError('The saved geometry is not ready for texturing.') + manifest = json.loads(_file(folder, 'geometry-manifest.json').read_text(encoding='utf-8')) + bindings = manifest.get('files') + handoff_files = _handoff_files(folder) + if manifest.get('schema_version') != 1 or not isinstance(bindings, dict) or set(bindings) != set(handoff_files): + raise ValueError('The saved geometry handoff is invalid.') + destination = Path(destination) + if destination.is_symlink() or destination.exists() and any(destination.iterdir()): + raise ValueError('Geometry must be restored into an empty local directory.') + destination.mkdir(parents=True, exist_ok=True) + for name in handoff_files: + source = _file(folder, name) + expected = bindings[name] + if (not isinstance(expected, dict) or source.stat().st_size != expected.get('size') + or _digest(source) != expected.get('sha256')): + raise ValueError('The saved geometry failed integrity validation.') + target = destination / name + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source, target) + if target.stat().st_size != expected['size'] or _digest(target) != expected['sha256']: + raise ValueError('The restored geometry failed integrity validation.') + return destination diff --git a/hf_space/gpu_mode.py b/hf_space/gpu_mode.py new file mode 100644 index 0000000000000000000000000000000000000000..f16ed40c1b07ea3bfd13d78552fde195ace79ee4 --- /dev/null +++ b/hf_space/gpu_mode.py @@ -0,0 +1,14 @@ +"""Hardware mode without importing Torch or starting a GPU runtime.""" +import os + + +def enabled(): + return os.environ.get('SPACES_ZERO_GPU', '').lower() in ('1', 'true', 't') + + +def pending_fields(operations): + if not enabled(): + return {} + pending = next((op for op in operations if op.get('status') == 'queued'), None) + return {'gpu_mode': 'zerogpu', 'gpu_operation_id': pending['id'] if pending else None, + 'gpu_operation_kind': pending['kind'] if pending else None} diff --git a/hf_space/history_jobs.py b/hf_space/history_jobs.py index e6c6239dd5ff1f34e96ec6ee20e6fed990bde2f8..db3afb56ec1798812ca4e60159ecd2f290f717ff 100644 --- a/hf_space/history_jobs.py +++ b/hf_space/history_jobs.py @@ -1,4 +1,4 @@ -"""One browser-independent FIFO worker for generation, export and recovery.""" +"""Durable operations: direct user requests on ZeroGPU, one FIFO on paid GPUs.""" from __future__ import annotations import logging @@ -11,6 +11,7 @@ LOG = logging.getLogger(__name__) FAILURE_MESSAGES = { "foreground": "Background removal failed. Please retry or choose another image.", "generate": "Generation failed. Please submit again.", + "texture": "Texture generation failed. Your geometry is saved; please retry texturing.", "export": "GLB export failed. Please retry.", "recover": "Preview recovery failed. Please retry.", } @@ -28,9 +29,17 @@ class Runner: self._thread = None self._start_lock = threading.Lock() self._pending_terminal = None + self._operation_guard = threading.Lock() + self._operation_locks = {} + self._pending_terminals = {} self.error = None def start(self): + from hf_space.gpu_mode import enabled + if enabled(): + # ZeroGPU work must retain the requesting visitor's Gradio context. + # Even an accidental call cannot create an anonymous FIFO worker. + return None with self._start_lock: if self._thread is not None: return self._thread @@ -63,10 +72,8 @@ class Runner: self.store.changed.wait(self.poll_interval) self.store.changed.clear() - def _finish_pending(self): - if self._pending_terminal is None: - return False - operation_id, kind, result, error = self._pending_terminal + def _persist_terminal(self, outcome): + operation_id, kind, result, error = outcome try: self.store.finish_operation(operation_id, result, error=error) except ValidationError: @@ -74,75 +81,127 @@ class Runner: # to rerun costly inference or block every later queued job. LOG.exception("Personal-history %s returned invalid output", kind) self.store.finish_operation(operation_id, error=FAILURE_MESSAGES[kind]) + + def _finish_pending(self): + if self._pending_terminal is None: + return False + self._persist_terminal(self._pending_terminal) self._pending_terminal = None return True - def run_once(self): - """Execute at most one durable operation; useful for deterministic tests.""" + def run_once(self, *, operation_id=None, owner=None): + """Run the selected visitor's operation, or advance the paid-GPU FIFO.""" + from hf_space.gpu_mode import enabled + if enabled(): + return self._run_selected(operation_id, owner) if not self.store.execution_lock.acquire(blocking=False): return False try: - if self._finish_pending(): + if self._finish_pending() and operation_id is None: return True - pending = self.store.next_operation() - if pending is None or self._stop.is_set(): - return False - operation = pending["operation"] - kind = operation["kind"] - if not self.readiness(kind): + pending = (self.store.get_operation(owner, operation_id) if operation_id is not None + else self.store.next_operation()) + outcome = self._execute_pending(pending, zero_gpu=False) + if outcome is None: return False - job = self.store.claim_operation(operation["id"]) - operation = next(item for item in job["operations"] if item["id"] == operation["id"]) - job["operation"] = operation - job["operation_id"] = operation["id"] - job["operation_kind"] = kind - watcher_stop = threading.Event() - update_lock = threading.Lock() - last_persisted = [0.0] - - def update(values=None, **changes): - if values is not None and changes: - raise ValidationError("Supply one progress object.") - value = values if values is not None else changes - if isinstance(value, dict) and set(value) == {"progress"}: - value = value["progress"] - # This callback can update progress only. It never edits owner, - # the task identity, persisted input, seed, or completion status. - with update_lock: - now = time.monotonic() - persist = now - last_persisted[0] >= 3.0 - self.store.update_progress(operation["id"], value, persist=persist) - if persist: - last_persisted[0] = now - - def watch(): - while not watcher_stop.wait(1.0): - try: - value = self.progress_reader(job, operation) - if value: - update(value) - except Exception: - # A transient progress read is not a failed GPU task. - LOG.warning("Could not read current task progress", exc_info=True) - - watcher = None - if self.progress_reader is not None: - watcher = threading.Thread(target=watch, name="personal-task-progress", daemon=True) - watcher.start() - result, failure = None, None - try: - result = self.handlers[kind](job, update) - except Exception: - failure = FAILURE_MESSAGES[kind] - LOG.exception("Personal-history %s task failed", kind) - finally: - watcher_stop.set() - if watcher is not None: - watcher.join(timeout=5) # Retain the outcome in memory if persistence temporarily fails. # Retry only the terminal write, never the completed GPU callback. - self._pending_terminal = (operation["id"], kind, result, failure) + self._pending_terminal = outcome self._finish_pending() return True finally: self.store.execution_lock.release() + + def _run_selected(self, operation_id, owner): + if operation_id is None: + return False + # Validate ownership before looking at an in-flight lock or retained + # result. A duplicate request never reveals another visitor's task. + self.store.get_operation(owner, operation_id) + with self._operation_guard: + lock = self._operation_locks.setdefault(operation_id, threading.Lock()) + if not lock.acquire(blocking=False): + return False + try: + pending = self.store.get_operation(owner, operation_id) + with self._operation_guard: + outcome = self._pending_terminals.get(operation_id) + if outcome is None: + outcome = self._execute_pending(pending, zero_gpu=True) + if outcome is None: + return False + with self._operation_guard: + self._pending_terminals[operation_id] = outcome + self._persist_terminal(outcome) + with self._operation_guard: + del self._pending_terminals[operation_id] + return True + finally: + # Registry changes and acquisition share this short lock, so a + # waiter cannot acquire an old lock while a new one is installed. + with self._operation_guard: + lock.release() + del self._operation_locks[operation_id] + + def _execute_pending(self, pending, *, zero_gpu): + if pending is None or self._stop.is_set(): + return None + operation = pending["operation"] + if operation['status'] != 'queued': + return None + kind = operation["kind"] + if not self.readiness(kind): + return None + job = self.store.claim_operation(operation["id"]) + operation = next(item for item in job["operations"] if item["id"] == operation["id"]) + job["operation"] = operation + job["operation_id"] = operation["id"] + job["operation_kind"] = kind + watcher_stop = threading.Event() + update_lock = threading.Lock() + last_persisted = [0.0] + + def update(values=None, **changes): + if values is not None and changes: + raise ValidationError("Supply one progress object.") + value = values if values is not None else changes + if isinstance(value, dict) and set(value) == {"progress"}: + value = value["progress"] + # This callback can update progress only. It never edits owner, + # the task identity, persisted input, seed, or completion status. + with update_lock: + now = time.monotonic() + persist = now - last_persisted[0] >= 3.0 + self.store.update_progress(operation["id"], value, persist=persist) + if persist: + last_persisted[0] = now + + def watch(): + while not watcher_stop.wait(1.0): + try: + value = self.progress_reader(job, operation) + if value: + update(value) + except Exception: + # A transient progress read is not a failed GPU task. + LOG.warning("Could not read current task progress", exc_info=True) + + watcher = None + if self.progress_reader is not None: + watcher = threading.Thread(target=watch, name="personal-task-progress", daemon=True) + watcher.start() + result, failure = None, None + try: + result = self.handlers[kind](job, update) + except Exception as error: + if zero_gpu: + from hf_space.zerogpu_jobs import public_gpu_error + failure = public_gpu_error(error, FAILURE_MESSAGES[kind]) + else: + failure = FAILURE_MESSAGES[kind] + LOG.exception("Personal-history %s task failed", kind) + finally: + watcher_stop.set() + if watcher is not None: + watcher.join(timeout=5) + return operation["id"], kind, result, failure diff --git a/hf_space/history_routes.py b/hf_space/history_routes.py index 39c1dca96fc1ad43855441eaeae6beaaa18b1bea..53ef7a78a97eb63d4bb131a4c754a24651a8420a 100644 --- a/hf_space/history_routes.py +++ b/hf_space/history_routes.py @@ -221,6 +221,7 @@ def install_routes( examples_provider: Callable | None = None, example_file_provider: Callable | None = None, prepare_example: Callable | None = None, + prepare_models: Callable | None = None, preview_provider: Callable | None = None, public_gradio_paths=(), trusted_origins=(), @@ -275,8 +276,9 @@ def install_routes( async def require_ready(): state = await invoke(deployment_status, _limiter=status_limiter) - if (state.get('native_build', {}).get('stage') != 'compiled' - or state.get('models', {}).get('stage') != 'ready_released'): + if (state.get('gpu_mode') != 'zerogpu' + and (state.get('native_build', {}).get('stage') != 'compiled' + or state.get('models', {}).get('stage') != 'ready_released')): raise HTTPException(503, 'Models are loading in the background. Please wait until loading is complete.') if not store.ready: raise HTTPException(503, 'Preparing your workspace. Please wait a moment.') @@ -374,6 +376,9 @@ def install_routes( result['status_url'] = f'{BASE}/uploads/{upload_id}' if upload['foreground'].get('status') == 'complete' and upload.get('files', {}).get('foreground.png'): result['foreground_url'] = f'{BASE}/uploads/{upload_id}/files/foreground.png' + for field in ('gpu_mode', 'gpu_operation_id', 'gpu_operation_kind'): + if field in upload: + result[field] = upload[field] return result def prime_completed_upload(owner, upload): @@ -487,6 +492,27 @@ def install_routes( state = await invoke(deployment_status, _limiter=status_limiter) return {**state, 'workspace': workspace_status()} + @router.post('/gpu/prepare') + async def prepare_gpu_task(request: Request): + identity = require_mutation(request, trusted_origins, csrf_secret) + available() + value = await body(request, ('operation_id',)) + pending = await invoke(store.get_operation, identity.owner, value.get('operation_id'), + _limiter=status_limiter) + state = await invoke(deployment_status, _limiter=status_limiter) + if state.get('gpu_mode') != 'zerogpu' or prepare_models is None: + raise HTTPException(404, 'On-demand preparation is unavailable.') + if pending['operation']['status'] == 'queued': + kind = pending['operation']['kind'] + required = ('foreground',) if kind == 'foreground' else ('native_build', 'models') + if any(state.get(key, {}).get('stage') == 'failed' for key in required): + raise HTTPException(503, 'Model preparation failed. Please try again later.') + # This callback only starts the shared CPU preparation controller; + # GPU allocation remains in the visitor's later Gradio request. + await invoke(prepare_models, kind, _limiter=mutation_limiter) + return {'operation_id': pending['operation']['id'], + 'status': pending['operation']['status'], 'gpu_mode': 'zerogpu'} + @router.get('/examples') async def examples(): # This manifest only describes committed public assets. It does not @@ -561,7 +587,7 @@ def install_routes( models = state.get('models', {}) if isinstance(state, dict) else {} if native.get('stage') == 'failed' or models.get('stage') == 'failed': raise HTTPException(503, 'Model loading failed. Please try again later.') - if native.get('stage') != 'compiled' or models.get('stage') != 'ready_released': + if state.get('gpu_mode') != 'zerogpu' and (native.get('stage') != 'compiled' or models.get('stage') != 'ready_released'): raise HTTPException(503, 'Models are loading in the background. Please wait until loading is complete.') upload_id = value.get('upload_id') if value.get('example'): @@ -637,7 +663,8 @@ def install_routes( await require_ready() available() value = await body(request, ('idempotency_key',)) - function = store.submit_export if kind == 'export' else store.submit_recover + function = {'export': store.submit_export, 'recover': store.submit_recover, + 'texture': store.submit_texture}[kind] job = await invoke(function, identity.owner, job_id, key(value.get('idempotency_key')), _limiter=mutation_limiter) return await invoke(snapshot, job) @@ -645,6 +672,10 @@ def install_routes( async def export(job_id: str, request: Request): return await operation(request, job_id, 'export') + @router.post('/jobs/{job_id}/texture') + async def texture(job_id: str, request: Request): + return await operation(request, job_id, 'texture') + @router.post('/jobs/{job_id}/recover') async def recover(job_id: str, request: Request): return await operation(request, job_id, 'recover') @@ -668,7 +699,8 @@ def install_routes( local, release = await invoke(prepare, _limiter=artifact_limiter if name.endswith('.glb') else image_limiter) except RuntimeError: raise HTTPException(503, 'Artifact storage is temporarily unavailable. Please retry.') from None - media = mimetypes.guess_type(name)[0] or 'application/octet-stream' + media = {'.webp': 'image/webp', '.png': 'image/png'}.get(Path(name).suffix.lower()) + media = media or mimetypes.guess_type(name)[0] or 'application/octet-stream' if name.lower().endswith('.glb'): media = 'model/gltf-binary' return PrivateFileResponse(local, media_type=media, filename=Path(name).name, diff --git a/hf_space/history_service.py b/hf_space/history_service.py index b4ca68b8e973d784860768eb3351dbb711eeb7b9..38ffb85ec410aea353e69f6ccb5d2fdced9fdd82 100644 --- a/hf_space/history_service.py +++ b/hf_space/history_service.py @@ -10,7 +10,7 @@ import logging from PIL import Image -from hf_space import release_backend, native_build, generation_progress, extraction_progress, ovoxel_results, example_assets +from hf_space import release_backend, native_build, generation_progress, extraction_progress, ovoxel_results, example_assets, geometry_results from hf_space.history_store import Store, NotFound from hf_space.history_jobs import Runner from hf_space.history_keepalive import WorkKeepAlive @@ -23,11 +23,22 @@ LOG = logging.getLogger(__name__) def deployment_status(): + from hf_space.gpu_mode import enabled native = native_build.status() models = release_backend.model_status() + native_steps = 7 # CUDA setup, five native wheels, dependency validation. + completed = (native_steps if native.get('stage') == 'compiled' else + 1 + len(native.get('completed', [])) if native.get('stage') == 'compiling' else 0) + native_progress = {'label': native.get('current') or 'Preparing native dependencies', + 'completed': min(native_steps, completed), 'total': native_steps, 'unit': 'steps'} + foreground = release_backend.foreground_status() if enabled() else None # A public readiness check must not expose traceback text, runtime file # paths, private configuration or job identifiers. - return {'native_build': {key: native[key] for key in ('stage', 'current') if key in native}, + return {'runtime_id': models.get('epoch'), + **({'gpu_mode': 'zerogpu', 'foreground': { + key: foreground[key] for key in ('stage', 'startup_progress') if key in foreground}} if enabled() else {}), + 'native_build': {**{key: native[key] for key in ('stage', 'current') if key in native}, + 'startup_progress': native_progress}, 'models': {key: models[key] for key in ('stage', 'startup_progress') if key in models}} @@ -36,7 +47,7 @@ class Service: self._boot_started = time.monotonic() self.store = store if store is not None else Store() self.private_files = PrivateFileCache() - self.runner = Runner(self.store, {'generate': self.generate, 'export': self.export, + self.runner = Runner(self.store, {'generate': self.generate, 'texture': self.texture, 'export': self.export, 'recover': self.recover, 'foreground': self.foreground}, ready=self.ready, progress=self.progress) self.keepalive = WorkKeepAlive(self.active_work) @@ -61,6 +72,18 @@ class Service: if self._started: return self._started = True + from hf_space.gpu_mode import enabled + if enabled(): + # History is available independently of models. Opening the page, + # polling status, and restoring saved records never prepare weights. + def initialize_history(): + try: + self.store.initialize() + except Exception: + LOG.exception('Could not initialize personal history') + threading.Thread(target=initialize_history, name='history-startup', daemon=True).start() + atexit.register(self.runner.stop) + return # One small read-only HTTP bootstrap lets existing history cards appear # during model loading. It never traverses the Bucket mount, writes task # state, triggers model preparation or waits on the weight controller. @@ -96,13 +119,24 @@ class Service: atexit.register(self.runner.stop) atexit.register(self.keepalive.stop) + @staticmethod + def prepare_models(kind='generate'): + """Start CPU preparation only after an explicit owner-checked task.""" + from hf_space.gpu_mode import enabled + if enabled(): + if kind == 'foreground': + release_backend.start_foreground_preparation() + else: + release_backend.start_preparation() + @staticmethod def ready(kind): state = deployment_status() + if state.get('gpu_mode') == 'zerogpu' and kind == 'foreground': + return state['foreground'].get('stage') == 'ready_released' native = state['native_build'].get('stage') == 'compiled' - # Startup constructs the shared runtime and opens a CUDA context outside - # the FIFO worker. Exports/recovery must wait for it too, avoiding GPU - # overlap and contention with the initial weight staging/loading. + # GPU execution requires the CPU pool and native extensions. Page, + # history and task admission have an independent workspace gate. return native and state['models'].get('stage') == 'ready_released' def generate(self, job, update): @@ -112,12 +146,31 @@ class Service: input_path = self.store.resolve_file(job['owner'], job['id'], 'foreground.png' if prepared else 'input.png') with Image.open(input_path) as source: image = source.convert('RGBA') - update(label='Generating geometry and reference materials') - mesh, _asset, report = release_backend.generate(image, int(job['seed']), generation_token=token, + split = job.get('workflow') == 'geometry_texture' + update(label='Generating geometry' if split else 'Generating geometry and reference materials') + generate = release_backend.generate_geometry if split else release_backend.generate + mesh, _asset, report = generate(image, int(job['seed']), generation_token=token, **({'foreground_prepared': True} if prepared else {})) - generation_progress.update(token, status='saving', label='Saving the 3D asset') - update(label='Saving the 3D asset') - identifier = ovoxel_results.persist(mesh, report) + return self._publish_generated(mesh, report, job, token, update, geometry=split) + + def texture(self, job, update): + """Resume from the saved white model; never regenerate its geometry.""" + import tempfile + if job.get('workflow') != 'geometry_texture' or not job.get('geometry_asset_id'): + raise ValueError('A saved geometry result is required before generating textures.') + token = job['operation_id'] + update(label='Loading saved geometry') + with tempfile.TemporaryDirectory(prefix='supratex-texture-input-') as temporary: + folder = geometry_results.restore(job['geometry_asset_id'], Path(temporary)) + mesh, _asset, report = release_backend.generate_texture( + folder, int(job['seed']), generation_token=token) + return self._publish_generated(mesh, report, job, token, update) + + def _publish_generated(self, mesh, report, job, token, update, *, geometry=False): + label = 'Saving geometry' if geometry else 'Saving the 3D asset' + generation_progress.update(token, status='saving', label=label) + update(label=label) + identifier = (geometry_results if geometry else ovoxel_results).persist(mesh, report) local = Path(mesh).parent files, metadata = ovoxel_results.local_manifest(local) # Prime the first actual view from local generation output. This is diff --git a/hf_space/history_store.py b/hf_space/history_store.py index 9fade79436cc83a01b6404017380d51104e61b4e..e3ddb6a133e6f6b716896eb70c0f1923c8f9e7d5 100644 --- a/hf_space/history_store.py +++ b/hf_space/history_store.py @@ -26,7 +26,8 @@ from PIL import Image, ImageOps, UnidentifiedImageError ID = re.compile(r"^[0-9a-f]{32}$") FRAME = re.compile(r"^[a-z][a-z0-9_]{0,47}-[0-9]{2}\.webp$") -KINDS = {"generate", "export", "recover", "foreground"} +KINDS = {"generate", "texture", "export", "recover", "foreground"} +GEOMETRY_TEXTURE = "geometry_texture" STATUSES = {"queued", "running", "complete", "failed", "interrupted"} MAX_UPLOAD_BYTES = 20 * 1024 * 1024 MAX_PIXELS = 24_000_000 @@ -391,6 +392,15 @@ class Store: if not upload: if record.get("status") not in STATUSES or not isinstance(record.get("operations"), list): raise ValidationError("Invalid job state.") + if record.get("workflow") not in (None, GEOMETRY_TEXTURE): + raise ValidationError("Invalid job workflow.") + if record.get("geometry_asset_id") is not None: + _id(record["geometry_asset_id"]) + if record.get("workflow") == GEOMETRY_TEXTURE: + if record.get("texture_status", "not_requested") not in STATUSES | {"not_requested"}: + raise ValidationError("Invalid texture state.") + if record["status"] == "complete" and not record.get("geometry_asset_id"): + raise ValidationError("The geometry checkpoint is unavailable.") for operation in record["operations"]: _id(operation.get("id")) if operation.get("kind") not in KINDS or operation.get("status") not in STATUSES: @@ -405,6 +415,8 @@ class Store: _key(alias) if (operation["kind"] == "foreground") != (record.get("purpose") == "foreground"): raise ValidationError("Invalid foreground task.") + if operation["kind"] == "texture" and record.get("workflow") != GEOMETRY_TEXTURE: + raise ValidationError("Invalid texture workflow.") if record.get("purpose") == "foreground" and record.get("upload_id") != record["id"]: raise ValidationError("Invalid foreground upload.") asset_id = record.get("asset_id") @@ -550,6 +562,8 @@ class Store: result["foreground"] = {"status": state["status"], "progress": state.get("progress", {})} if state.get("error"): result["foreground"]["error"] = state["error"] + from hf_space.gpu_mode import pending_fields + result.update(pending_fields(foreground['operations'])) if state["status"] == "complete" and state.get("foreground_metadata"): result.setdefault("files", {})["foreground.png"] = deepcopy(state["files"]["foreground.png"]) return result @@ -685,10 +699,12 @@ class Store: "idempotency_key": key, "enqueue_order": sequence, "created_at": now, "updated_at": now, "progress": {}, "error": None} - def submit(self, owner, upload_id, seed, idempotency_key): + def submit(self, owner, upload_id, seed, idempotency_key, *, workflow=GEOMETRY_TEXTURE): self._require_ready() _owner(owner) _key(idempotency_key) + if workflow not in (None, GEOMETRY_TEXTURE): + raise ValidationError("Invalid generation workflow.") if type(seed) is not int or not 0 <= seed <= 2**31 - 1: raise ValidationError("Seed must be an integer between 0 and 2147483647.") with self._durable_mutation(("admission", owner)): @@ -703,9 +719,7 @@ class Store: if not same_input or job["seed"] != seed: raise Conflict("This request key already belongs to another submission.") return self._snapshot(job) - pending = sum(job["owner"] == owner and job["status"] in ("queued", "running") - and job.get("purpose") != "foreground" - for job in self._jobs.values()) + pending = self._pending_generations(owner) if pending >= self.max_pending_per_owner: raise LimitExceeded("You can have at most three unfinished generations.") foreground = self._jobs.get(upload_id) @@ -741,6 +755,11 @@ class Store: "error": None, "export_status": "not_requested", "export_progress": {}, "export_error": None, "recover_status": "not_requested", "recover_progress": {}, "operations": [operation], "files": deepcopy(upload.get("files", {}))} + # Operator-managed background refreshes explicitly retain the full + # pipeline. HTTP submissions never accept a workflow override. + if workflow == GEOMETRY_TEXTURE: + job.update(workflow=workflow, geometry_asset_id=None, texture_status="not_requested", + texture_progress={}, texture_error=None) if foreground: job.update(foreground_prepared=True, foreground_metadata=deepcopy(foreground["foreground_metadata"])) job["files"]["foreground.png"] = deepcopy(foreground["files"]["foreground.png"]) @@ -767,10 +786,18 @@ class Store: self._revision += 1 self._job_revisions[job_id] = self._revision + def _pending_generations(self, owner): + return sum(job["owner"] == owner and any( + op["kind"] in ("generate", "texture") and op["status"] in ("queued", "running") + for op in job["operations"]) for job in self._jobs.values()) + def _submit_operation(self, kind, owner, job_id, idempotency_key): self._require_ready() _key(idempotency_key) - with self._durable_mutation(("job", job_id)): + # Share admission ordering with new generations: a pending texture also + # occupies an unfinished-generation slot, even while its white preview + # remains complete and viewable. + with self._writer_for(("admission", owner)), self._durable_mutation(("job", job_id)): with self._lock: original = deepcopy(self._authorized(self._jobs, owner, job_id)) previous = self._keys.get((owner, kind, idempotency_key)) @@ -780,10 +807,18 @@ class Store: return self._snapshot(original) if original["status"] != "complete" or not original.get("asset_id"): raise Conflict("Wait for generation to finish before requesting this operation.") + if kind == "texture": + if (original.get("workflow") != GEOMETRY_TEXTURE + or not isinstance(original.get("geometry_asset_id"), str) + or not ID.fullmatch(original["geometry_asset_id"])): + raise Conflict("This generation has no separate geometry checkpoint to texture.") + elif kind == "export" and original.get("workflow") == GEOMETRY_TEXTURE: + if original.get("texture_status") != "complete": + raise Conflict("Generate the texture before exporting this model.") existing = next((index for index in range(len(original["operations"]) - 1, -1, -1) if original["operations"][index]["kind"] == kind and (original["operations"][index]["status"] in ("queued", "running") - or kind == "export" and original["operations"][index]["status"] == "complete")), + or kind in ("export", "texture") and original["operations"][index]["status"] == "complete")), None) if existing is not None: # Persist merged request keys too: retrying a second click @@ -793,6 +828,13 @@ class Store: self._commit(job) with self._lock: return self._snapshot(job) + if any(op["status"] in ("queued", "running") and + (kind == "texture" or op["kind"] == "texture") for op in original["operations"]): + raise Conflict("Wait for the current model operation to finish before requesting another stage.") + if kind == "texture": + with self._lock: + if self._pending_generations(owner) >= self.max_pending_per_owner: + raise LimitExceeded("You can have at most three unfinished generations.") job = deepcopy(original) operation = self._new_operation(kind, idempotency_key) job["operations"].append(operation) @@ -804,6 +846,9 @@ class Store: def submit_export(self, owner, job_id, idempotency_key): return self._submit_operation("export", owner, job_id, idempotency_key) + def submit_texture(self, owner, job_id, idempotency_key): + return self._submit_operation("texture", owner, job_id, idempotency_key) + def submit_recover(self, owner, job_id, idempotency_key): return self._submit_operation("recover", owner, job_id, idempotency_key) @@ -821,6 +866,11 @@ class Store: if operation["status"] == "queued"), key=lambda entry: entry[0]) def _queue_state(self): + from hf_space.gpu_mode import enabled + if enabled(): + # ZeroGPU requests run independently; a global FIFO rank would be + # misleading. The HF scheduler alone controls GPU availability. + return {} running = sum(operation["status"] == "running" for record in self._jobs.values() for operation in record["operations"]) positions = {} @@ -843,9 +893,11 @@ class Store: keys = ("id", "upload_id", "status", "seed", "created_at", "updated_at", "asset_id", "progress", "error", "export_status", "export_progress", "export_error", "export_postprocess", "recover_status", "recover_progress", "recover_error", "preview", - "foreground_prepared") + "foreground_prepared", "workflow", "texture_status", "texture_progress", "texture_error") result = {key: deepcopy(job[key]) for key in keys if key in job} result["display_ready"] = "display.webp" in job.get("files", {}) + from hf_space.gpu_mode import pending_fields + result.update(pending_fields(job.get('operations', []))) return result def get_job(self, owner, job_id): @@ -929,6 +981,16 @@ class Store: _, identifier, operation = queued[0] return {"job": self._snapshot(self._jobs[identifier]), "operation": deepcopy(operation)} + def get_operation(self, owner, operation_id): + """Look up an operation only for its verified owner, including cutouts.""" + self._require_ready() + _owner(owner) + with self._lock: + job, index = self._operation(operation_id) + if job['owner'] != owner or job.get('deleted_at'): + raise NotFound('The task is not available.') + return {'job': self._snapshot(job), 'operation': deepcopy(job['operations'][index])} + def has_pending(self): """Whether this initialized process has any queued or running work.""" with self._lock: @@ -1017,14 +1079,20 @@ class Store: if operation["kind"] != "export": raise ValidationError("Only GLB exports can report texture postprocessing.") job["export_postprocess"] = _export_postprocess(result["export_postprocess"]) - if operation["kind"] == "generate": + replacing_preview = operation["kind"] == "texture" + if operation["kind"] in ("generate", "texture"): try: asset_id = _id(result.get("asset_id")) except NotFound as invalid: raise ValidationError("The generated asset identifier is invalid.") from invalid if not (self.asset_root / asset_id / "ready.json").is_file(): raise ValidationError("The generated asset has not finished saving.") + if replacing_preview and (job.get("workflow") != GEOMETRY_TEXTURE + or asset_id == job.get("geometry_asset_id")): + raise ValidationError("Texture output must preserve the saved geometry checkpoint.") job["asset_id"] = asset_id + if operation["kind"] == "generate" and job.get("workflow") == GEOMETRY_TEXTURE: + job["geometry_asset_id"] = asset_id elif "asset_id" in result and result["asset_id"] != job["asset_id"]: raise ValidationError("A task cannot replace another asset.") if operation["kind"] == "export" and not (self.asset_root / job["asset_id"] / "asset.glb").is_file(): @@ -1032,7 +1100,7 @@ class Store: metadata = result.get("preview_metadata") if metadata is not None: metadata = _preview_metadata(metadata) - elif job.get("preview_metadata"): + elif job.get("preview_metadata") and not replacing_preview: metadata = _preview_metadata(job["preview_metadata"]) elif job.get("asset_id") and (self.asset_root / job["asset_id"] / "preview.json").is_file(): try: @@ -1040,6 +1108,8 @@ class Store: except ValueError as invalid: raise ValidationError("The saved preview metadata is invalid.") from invalid metadata = _preview_metadata(metadata) + if replacing_preview and metadata is None: + raise ValidationError("The textured preview metadata is unavailable.") if "preview" in result: job["preview"] = _preview(result["preview"]) elif metadata is not None: @@ -1056,6 +1126,15 @@ class Store: # handlers report their local sizes while publishing. manifest = self._capture_asset_files(job, metadata) self._validate_asset_manifest(job, manifest, metadata or job.get("preview_metadata")) + if replacing_preview: + frames = {f"frames/{mode}-{view:02d}.webp" for mode in metadata["modes"] + for view in range(metadata["views"])} + if not frames.issubset(manifest): + raise ValidationError("The textured preview has not finished saving.") + # The asset changes atomically with its descriptors. Keep + # the owner's saved inputs; remove only obsolete model files. + job["files"] = {name: value for name, value in job.get("files", {}).items() + if name in PRIVATE_IMAGE_NAMES} job.setdefault("files", {}).update(manifest) elapsed = operation.get("progress", {}).get("elapsed_seconds") operation.update(status="complete", error=None) diff --git a/hf_space/history_ui.css b/hf_space/history_ui.css index 2a065cf8b458064cf6a4657252c63cece418e3fb..358077bfbbd129560666e2f372ddd67f4f2d4dbb 100644 --- a/hf_space/history_ui.css +++ b/hf_space/history_ui.css @@ -7,15 +7,25 @@ #st-history-header .st-history-heading p{font-size:12px;color:#9c9c9a;margin:5px 0 0;line-height:1.5} #st-resource-footer{width:auto!important;min-width:0!important;flex:none!important;margin:10px clamp(12px,1.2vw,24px) 12px!important;padding:0!important;border:0!important;background:transparent!important} #st-resource-footer .html-container{padding:0!important;border:0!important;background:transparent!important} -#st-resource-footer .st-resource-notice{padding:9px 12px;border:1px solid #655536;border-left:3px solid #edc96b;border-radius:7px;background:#29251c;color:#e5d5ac;font-size:12px;line-height:1.5;font-family:Arial,sans-serif} -#st-resource-footer .st-resource-notice strong{color:#f3d386;font-weight:600} +#st-resource-footer .st-resource-notice{padding:9px 12px;border:1px solid #655536;border-left:3px solid #edc96b;border-radius:7px;background:#29251c;color:#e5d5ac;font-size:12px;line-height:1.8;font-family:Arial,sans-serif} +#st-resource-footer strong{color:#f3d386;font-weight:600} +#st-resource-footer .st-plan-link{display:inline-flex;align-items:center;gap:4px;vertical-align:middle;white-space:nowrap;color:inherit;text-decoration:none} +#st-resource-footer .st-plan-link:hover{text-decoration:underline} +#st-resource-footer .st-plan-link:focus-visible{outline:2px solid #f2cd70;outline-offset:3px;border-radius:4px} +#st-resource-footer .st-hf-logo{display:inline-block;width:18px;height:17px;object-fit:contain} +#st-resource-footer .st-hf-free{display:inline-block;border:1px solid #ffe4a0;border-radius:6px;padding:0 7px;background:#f2cd70;color:#241f13;font-size:12px;font-weight:800;line-height:20px;letter-spacing:.04em} +/* HF LabelPro: official dark-theme colors and badge geometry from huggingface.co/pricing. */ +#st-resource-footer .st-hf-pro{display:inline-block;transform:skewX(-12deg);border:1px solid #e5e7eb;border-radius:6px;padding:0 6px;background:linear-gradient(to bottom right,oklch(65.6% .241 354.308) 0%,oklch(69.6% .17 162.48) 50%,oklch(76.9% .188 70.08) 100%);color:#000;font-size:12px;font-weight:700;line-height:20px;box-shadow:0 10px 15px -3px oklch(69.6% .17 162.48 / .2),0 4px 6px -4px oklch(69.6% .17 162.48 / .2)} #st-oauth-bootstrap{display:none!important} #st-app.st-history-app{height:var(--st-workspace-height,calc(var(--st-available-height,100dvh) - 190px));width:100%;padding:0 clamp(12px,1.2vw,24px);margin:0;background:transparent;font-family:Arial,sans-serif} #st-app.st-history-app [hidden]{display:none!important} #st-app.st-history-app #st-workspace{position:relative;display:grid!important;grid-template-columns:minmax(260px,1fr) minmax(260px,1fr) clamp(112px,9vw,154px) clamp(112px,9vw,154px);gap:12px!important;height:100%;margin:0;min-height:0;align-items:stretch!important} #st-app.st-history-app .st-native-panel{display:flex!important;flex-direction:column!important;gap:9px!important;border:1px solid #373735;border-radius:15px;background:#191919;padding:12px!important;height:100%;min-width:0;min-height:0;overflow:hidden;box-shadow:none} #st-app.st-history-app #st-input-panel{grid-column:1;overflow:auto} -#st-app.st-history-app #st-output-panel{grid-column:2;overflow:hidden} +#st-app.st-history-app #st-input-panel[data-busy=true]{pointer-events:none;border-color:#494337} +#st-app.st-history-app #st-input-panel[data-busy=true] :is(button,input,summary){opacity:.46;cursor:not-allowed} +#st-app.st-history-app #st-examples button:disabled{opacity:.46;cursor:not-allowed} +#st-app.st-history-app #st-output-panel{grid-column:2;overflow:auto} #st-app.st-history-app #st-history-panel{grid-column:3;contain:none} #st-app.st-history-app #st-examples-panel{grid-column:4;contain:none} #st-app.st-history-app h2,#st-app.st-history-app .st-panel-title{font-size:12px;font-weight:400;line-height:1.4;margin:0;color:#ccc;flex:0 0 auto} @@ -31,7 +41,6 @@ #st-app.st-history-app .st-upload-actions button{border:0;background:none;color:#dfc775;font-size:11px;padding:3px} #st-app.st-history-app #st-upload-status{flex:none;display:flex;align-items:center;gap:8px;flex-wrap:wrap;font-size:11px;color:#c3bcaa;line-height:1.5} #st-app.st-history-app #st-upload-label{flex:1 1 180px} -#st-app.st-history-app #st-upload-meter{width:100%;height:4px;accent-color:#f2cd70;order:3} #st-app.st-history-app #st-upload-retry{border:1px solid #62583c;border-radius:5px;background:#28251e;color:#f2cd70;padding:4px 10px} #st-app.st-history-app #st-upload-status[data-state=complete]{color:#bcd5a8} #st-app.st-history-app #st-upload-status[data-state=failed]{color:#e0b6a7} @@ -48,12 +57,6 @@ #st-app.st-history-app #st-account-help{font-size:11px;line-height:1.6;color:#c3bcaa;margin:0} #st-app.st-history-app #st-cancel-signin{align-self:center;border:0;background:none;color:#c3bcaa;font-size:11px;padding:3px;text-decoration:underline;text-underline-offset:2px} #st-app.st-history-app #st-generate-control[data-notice=true] #st-gate-note{color:#f2cd70} -#st-app.st-history-app #st-native-startup{padding:9px 11px;border:1px solid #494337;border-radius:8px;background:#211f1b} -#st-app.st-history-app #st-native-startup strong{font-size:11px;font-weight:500;color:#e8d8ac} -#st-app.st-history-app #st-native-startup small{display:block;font-size:10px;color:#aaa;overflow-wrap:anywhere} -#st-app.st-history-app #st-native-startup progress{display:block;width:100%;height:4px;margin:7px 0;accent-color:#f2cd70} -#st-app.st-history-app .st-startup-heading{display:flex;justify-content:space-between;gap:7px;flex-wrap:wrap} -#st-app.st-history-app #st-startup-detail{font-size:10px;color:#bdb5a2} #st-app.st-history-app #st-viewer{position:relative;flex:1 1 0!important;display:flex;min-height:180px!important;overflow:hidden} #st-app.st-history-app .st-native-empty{display:grid;place-items:center;flex:1;width:100%;min-height:0;text-align:center;color:#999;font-size:13px;padding:20px} #st-app.st-history-app #st-model,#st-app.st-history-app #st-ovoxel{flex:1 1 0!important;width:100%;min-height:0!important;height:100%!important} @@ -68,6 +71,26 @@ #st-app.st-history-app #st-queue-history-link:focus-visible{outline:2px solid #f2cd70;outline-offset:3px;border-radius:2px} #st-app.st-history-app #st-job-progress progress{width:min(260px,80%);height:5px;accent-color:#f2cd70} #st-app.st-history-app #st-job-progress small{font-size:11px;color:#999} +#st-app.st-history-app #st-job-progress[data-resumable=true]{gap:8px;padding:10px 22px 34px} +#st-app.st-history-app #st-resume-gpu{appearance:none;flex:none;min-height:44px;max-width:100%;margin:4px 0 0;padding:10px 18px;border:1px solid #ffe4a0;border-radius:9px;background:#f2cd70;color:#241f13;font-size:13px;font-weight:700;line-height:1.4;cursor:pointer} +#st-app.st-history-app #st-resume-gpu:hover{background:#ffe09a;border-color:#ffedbc} +#st-app.st-history-app #st-resume-gpu:focus-visible{outline:2px solid #ffe4a0;outline-offset:3px} +#st-app.st-history-app #st-texture-actions{flex:none;display:flex;flex-direction:column;align-items:center;gap:8px;padding:10px 12px;border:1px solid #494337;border-radius:9px;background:#211f1b;text-align:center} +#st-app.st-history-app #st-texture-note,#st-app.st-history-app #st-texture-detail{font-size:12px;line-height:1.5;color:#c8c1b1;margin:0} +#st-app.st-history-app #st-texture{appearance:none;min-height:44px;width:100%;padding:10px 18px;border:1px solid #ffe4a0;border-radius:9px;background:#f2cd70;color:#241f13;font-size:13px;font-weight:700;cursor:pointer} +#st-app.st-history-app #st-texture:hover{background:#ffe09a} +#st-app.st-history-app #st-texture:focus-visible{outline:2px solid #ffe4a0;outline-offset:3px} +#st-app.st-history-app #st-texture:disabled{opacity:.55;cursor:default} +#st-app.st-history-app #st-texture-progress{width:100%;display:flex;flex-direction:column;align-items:center;gap:7px} +#st-app.st-history-app #st-texture-title{font-size:13px;font-weight:600;color:#efe3bf} +#st-app.st-history-app #st-texture-progress progress{width:min(260px,90%);height:5px;accent-color:#f2cd70} +#st-app.st-history-app #st-texture-time{font-size:11px;color:#999} +#st-app.st-history-app #st-output-panel[data-texture-progress=true] #st-viewer{display:none} +#st-app.st-history-app #st-output-panel[data-texture-progress=true] #st-texture-actions{flex:1 1 0;min-height:180px;justify-content:center;padding:22px;border:0;background:transparent;gap:12px} +#st-app.st-history-app #st-output-panel[data-texture-progress=true] #st-texture-progress{gap:12px} +#st-app.st-history-app #st-output-panel[data-texture-progress=true] #st-texture-title{font-size:18px;font-weight:500} +#st-app.st-history-app #st-output-panel[data-texture-progress=true] #st-texture-detail{font-size:13px;line-height:1.6;color:#bbb} +#st-app.st-history-app #st-output-panel[data-texture-progress=true] #st-texture{width:min(280px,100%);margin-top:4px} #st-app.st-history-app #st-export-note{font-size:11px;line-height:1.5;margin:0;flex:none;min-height:17px} #st-app.st-history-app #st-extract[data-action=download]{background:#e2c478!important;color:#241e10!important} #st-app.st-history-app #st-export-progress{flex:none;padding:10px 12px;border:1px solid #494337;border-radius:8px;background:#211f1b} @@ -130,6 +153,7 @@ #st-app.st-history-app :is(#st-input-panel,#st-output-panel,#st-history-panel,#st-examples-panel){grid-column:1;height:auto;min-width:0} #st-app.st-history-app #st-upload{height:350px!important;flex:none!important} #st-app.st-history-app #st-viewer{height:420px;flex:none!important} + #st-app.st-history-app #st-output-panel[data-texture-progress=true] #st-texture-actions{flex:none;min-height:420px} #st-app.st-history-app #st-history-panel,#st-app.st-history-app #st-examples-panel{height:300px;max-height:none} #st-app.st-history-app #st-history-list,#st-app.st-history-app #st-examples{grid-template-columns:repeat(auto-fill,minmax(95px,1fr))} #st-app.st-history-app #st-history-panel.expanded{inset:12px!important;width:calc(100vw - 24px)!important;height:calc(100dvh - 24px)!important;max-height:none!important} diff --git a/hf_space/history_ui.js b/hf_space/history_ui.js index 7c52952078991aab15a723b93b6ca6615a5576ec..bfea67e7197e4f55cbe3afd09bd3203f050d2f79 100644 --- a/hf_space/history_ui.js +++ b/hf_space/history_ui.js @@ -4,30 +4,121 @@ root.dataset.historyInstalled = 'true'; const $ = id => document.getElementById(id); const pending = job => ['queued', 'running'].includes(job?.status) || - ['queued', 'running'].includes(job?.export_status) || ['queued', 'running'].includes(job?.recover_status); + ['queued', 'running'].includes(job?.texture_status) || ['queued', 'running'].includes(job?.export_status) || ['queued', 'running'].includes(job?.recover_status); + const stagedJob = job => job?.workflow === 'geometry_texture'; + const texturesReady = job => !stagedJob(job) || job.texture_status === 'complete'; const terminal = state => ['native_build', 'models', 'workspace'].some(key => /^(failed|blocked)/.test(state?.[key]?.stage || '')); - const modelsReady = state => state?.native_build?.stage === 'compiled' && state?.models?.stage === 'ready_released'; - const ready = state => modelsReady(state) && (!state.workspace || state.workspace.stage === 'ready'); + const zeroGPU = state => state?.gpu_mode === 'zerogpu'; + const modelsReady = window.SupraZeroGPU.modelsReady; + const ready = window.SupraZeroGPU.workspaceReady; + const {foregroundOperation, preparationState, operationReady, preparationError, formatProgress} = window.SupraZeroGPU; const S = {identity: null, identityVersion: 0, deployment: null, jobs: new Map(), examples: new Map(), cursor: null, revision: null, historyLoaded: false, examplesLoaded: false, examplesPromise: null, historyError: '', historyFailures: 0, historyFailureSince: null, pageLoading: false, selected: null, selectionVersion: 0, previewLoaded: null, inputVersion: 0, source: null, uploadBusy: false, uploadState: null, submitBusy: false, expanded: false, jobsTimer: null, startupTimer: null, noticeTimer: null, exportTimer: null, previewController: null, stopped: false, - refreshPromise: null, statusPromise: null, lastIdentityCheck: 0, submitKey: null, + refreshPromise: null, statusPromise: null, lastIdentityCheck: 0, submitKey: null, resumeOperation: null, + textureBusy: null, textureKeys: new Map(), textureErrors: new Map(), localURL: null, silentAttempted: false, silentFrame: null, silentTimer: null, authPopup: null, authTimer: null, authBootstrapTimer: null, authFlow: null, authVersion: 0, authBusy: false, authFallback: false, identityPromise: null, authChecking: true, autoAuthAttempted: false, authMessage: ''}; + const gpuDispatcher = window.SupraZeroGPU.createDispatcher({ + identity: () => S.identity, version: () => S.identityVersion, stopped: () => S.stopped, + ready: record => operationReady(S.deployment, record), + prepare: (record, signal) => api('/gpu/prepare', {method: 'POST', json: {operation_id: record.gpu_operation_id}, signal}), + onPending() { void refreshStatus(); renderSelected(); }, + onError(record, message) { + const state = S.uploadState; + if (state && currentUpload(state) && foregroundOperation(record) && state.uploadId === record.upload_id) { + S.uploadBusy = false; state.retryAction = 'gpu'; state.gpuRecord = record; + clearTimeout(state.timer); uploadLabel(state, 'failed', message, true); + void refreshUpload(state); + } else notice(message); + renderSelected(); void refreshJobs(); + }, + onComplete(record) { + const state = S.uploadState; + if (state && currentUpload(state) && foregroundOperation(record) && state.uploadId === record.upload_id) void refreshUpload(state); + void refreshJobs(); + } + }); + const resumeGPU = document.createElement('button'); + resumeGPU.id = 'st-resume-gpu'; resumeGPU.type = 'button'; resumeGPU.hidden = true; + resumeGPU.textContent = 'Resume task'; + resumeGPU.addEventListener('click', async () => { + const job = S.selected?.kind === 'job' && S.jobs.get(S.selected.id); + if (!job || !S.identity?.authenticated || S.resumeOperation) return; + const identityVersion = S.identityVersion, operation = job.gpu_operation_id; + S.resumeOperation = operation; renderSelected(); + try { + // Reconnect to the current instance before using its readiness/progress. + const deployment = await refreshStatus(); + if (identityVersion !== S.identityVersion || S.stopped || !S.identity?.authenticated) return; + const latest = S.jobs.get(job.id); + if (!deployment || !ready(deployment)) { + notice('The service is reconnecting. Please try Resume task again.'); return; + } + if (latest?.gpu_operation_id === operation) void gpuDispatcher.retry(latest); + } finally { + S.resumeOperation = null; renderSelected(); + } + }); + $('st-job-progress').append(resumeGPU); + const needsGPUResume = record => record?.gpu_mode === 'zerogpu' && record.gpu_operation_id && + (!gpuDispatcher.requested(record.gpu_operation_id) || !!gpuDispatcher.failure(record.gpu_operation_id)); + const taskPreparing = record => record?.gpu_mode === 'zerogpu' && record.gpu_operation_id && + gpuDispatcher.requested(record.gpu_operation_id) && !gpuDispatcher.failure(record.gpu_operation_id) && !operationReady(S.deployment, record); + const startupBlocked = () => zeroGPU(S.deployment) ? /^(failed|blocked)/.test(S.deployment.workspace?.stage || '') : terminal(S.deployment); + const pollStartup = () => !startupBlocked() && (!ready(S.deployment) || gpuDispatcher.pending(record => + !operationReady(S.deployment, record) && !preparationError(S.deployment, record))); + function preparationLabel(record) { + const state = preparationState(S.deployment, record); + if (!foregroundOperation(record) && state === S.deployment?.native_build) return 'Preparing environment'; + const subject = foregroundOperation(record) ? 'background-removal model' : 'model weights'; + return (state.stage === 'staging_local_weights' ? 'Downloading ' : state.stage === 'loading_models' ? 'Loading ' : 'Preparing ') + subject; + } + function preparationText(record) { + const progress = preparationState(S.deployment, record).startup_progress || {}; + return preparationLabel(record) + ' · ' + formatProgress(progress).percent + '%'; + } + function preparationMeter(meter, record) { + const progress = preparationState(S.deployment, record).startup_progress || {}; + meter.max = 100; + if (taskPreparing(record)) meter.value = formatProgress(progress).percent; + else meter.removeAttribute('value'); + } // Gradio can finish its independent script callbacks in either order. Give // a late viewer installation only the current cache scope, never credentials. window.SupraPrivateAssetAccount = () => !S.stopped && S.identity?.authenticated ? String(S.identity.account_key || S.identity.username || '') : ''; const statusNames = {queued: 'Queued', running: 'Generating', complete: 'Generation Complete', failed: 'Failed', interrupted: 'Interrupted'}; + const zeroGPURecord = record => record?.gpu_mode === 'zerogpu' || !record?.gpu_mode && zeroGPU(S.deployment); + const leftInputBusy = () => S.submitBusy || !!S.textureBusy || !!S.resumeOperation || [...S.jobs.values()].some(job => { + const states = [job.status, job.texture_status]; + if (states.includes('running')) return true; + return states.includes('queued') && (!zeroGPURecord(job) || !!job.gpu_operation_id && + gpuDispatcher.requested(job.gpu_operation_id) && !gpuDispatcher.failure(job.gpu_operation_id)); + }); + const selectedGeometry = () => { + const job = S.selected?.kind === 'job' && S.jobs.get(S.selected.id); + return stagedJob(job) && job.status === 'complete' && S.source?.upload_id && + S.source.upload_id === job.upload_id && Math.round(Number($('st-seed').value)) === Number(job.seed) ? job : null; + }; + const textureFailure = job => (!['running', 'complete'].includes(job.texture_status) && S.textureErrors.get(job.id)) || + (job.gpu_operation_kind === 'texture' && gpuDispatcher.failure(job.gpu_operation_id)); + const textureProgressVisible = job => stagedJob(job) && job.status === 'complete' && !texturesReady(job) && + (S.textureBusy === job.id || ['queued', 'running'].includes(job.texture_status) && !textureFailure(job)); + const generationStatusName = job => stagedJob(job) ? ({queued: 'Preparing geometry', running: 'Generating geometry', + complete: 'Geometry ready', failed: 'Geometry failed', interrupted: 'Geometry interrupted'})[job.status] : + job.status === 'queued' && zeroGPURecord(job) ? 'Preparing' : statusNames[job.status]; const text = (id, value) => { $(id).textContent = value || ''; }; const show = (id, visible) => { $(id).hidden = !visible; }; const englishMessage = (value, fallback) => typeof value === 'string' && value.length <= 500 && !/\p{Script=Han}/u.test(value) ? value : fallback; + const generationProgressLabel = (value, fallback) => englishMessage( + value === 'Running the latest single-view pipeline' ? 'Starting generation…' : value, fallback); const loadingCardMessage = 'Please wait for the models to finish loading.'; const workspaceMessage = 'Preparing your workspace. Please wait a moment.'; let cardTip = null, cardTipTarget = null; @@ -37,9 +128,9 @@ } function gateCard(card) { const loading = !ready(S.deployment); - card.disabled = !S.identity?.authenticated || loading; + card.disabled = !S.identity?.authenticated || loading || card.dataset.example !== undefined && leftInputBusy(); if (S.identity?.authenticated && loading && !terminal(S.deployment)) - card.title = modelsReady(S.deployment) ? workspaceMessage : loadingCardMessage; + card.title = zeroGPU(S.deployment) || modelsReady(S.deployment) ? workspaceMessage : loadingCardMessage; else card.removeAttribute('title'); } function showCardTip(event) { @@ -330,6 +421,9 @@ updateGate(); syncAccount(); } function clearPreview() { + $('st-output-panel').dataset.textureProgress = 'false'; + show('st-texture-actions', false); + show('st-extract', true); show('st-export-note', true); show('st-view-switch', false); clearTimeout(S.exportTimer); S.exportTimer = null; $('st-model').replaceChildren(); $('st-ovoxel').replaceChildren(); @@ -343,6 +437,7 @@ function clearPrivate() { $('st-delete-dialog').close(); deleteTarget = null; deletedHistory.clear(); ++S.identityVersion; ++S.selectionVersion; ++S.inputVersion; + gpuDispatcher.reset(); clearUploadPreview(); clearHistoryThumbnails(); S.previewController?.abort(); S.previewController = null; @@ -351,6 +446,7 @@ S.historyFailures = 0; S.historyFailureSince = null; S.selected = null; S.previewLoaded = null; S.refreshPromise = null; S.pageLoading = false; S.source = null; S.uploadBusy = false; S.submitBusy = false; S.submitKey = null; + S.textureBusy = null; S.textureKeys.clear(); S.textureErrors.clear(); if (S.localURL) URL.revokeObjectURL(S.localURL); S.localURL = null; clearTimeout(S.jobsTimer); S.jobsTimer = null; @@ -397,6 +493,12 @@ error.status = response.status; error.code = data.code; throw error; } + const textureJob = method === 'POST' && /^\/jobs\/([^/]+)\/texture$/.exec(path); + if (textureJob && data?.id !== decodeURIComponent(textureJob[1])) + throw new Error('The texture response did not match this saved model. Please try again.'); + // Reading history never starts CPU preparation or spends GPU quota. Only + // the user's creation/retry POST establishes intent for this page session. + gpuDispatcher.observe(data, {requested: method === 'POST' && path !== '/gpu/prepare'}); return data; } async function refreshIdentity() { @@ -441,8 +543,17 @@ updateGate(); } function updateGate() { - const signedIn = !!S.identity?.authenticated, loading = !ready(S.deployment), failed = terminal(S.deployment); - const workspaceLoading = modelsReady(S.deployment) && loading; + const signedIn = !!S.identity?.authenticated, loading = !ready(S.deployment); + const inputBusy = leftInputBusy(), geometry = selectedGeometry(); + $('st-input-panel').inert = inputBusy; + $('st-input-panel').dataset.busy = String(inputBusy); + $('st-input-panel').setAttribute('aria-disabled', String(inputBusy)); + $('st-settings').inert = inputBusy; + $('st-settings').setAttribute('aria-disabled', String(inputBusy)); + const imageFlow = S.uploadBusy || !!(S.uploadState && currentUpload(S.uploadState)) || + !!(S.source?.foreground_required && !S.source.foreground_ready); + const failed = zeroGPU(S.deployment) ? /^(failed|blocked)/.test(S.deployment.workspace?.stage || '') : terminal(S.deployment); + const workspaceLoading = (zeroGPU(S.deployment) || modelsReady(S.deployment)) && loading; if (!loading && $('st-native-notice').textContent.startsWith('Models are loading in the background.')) { clearTimeout(S.noticeTimer); show('st-native-notice', false); } @@ -452,46 +563,29 @@ 'Model loading failed. Please refresh and try again later.'; else if (loading) reason = workspaceLoading ? workspaceMessage : 'Models are loading in the background. You can generate once loading is complete.'; else if (S.authBusy) reason = 'Connecting your Huggingface account…'; - else if (S.uploadBusy || S.source?.foreground_required && !S.source.foreground_ready) - reason = $('st-upload-label').textContent || 'Removing background…'; + else if (S.uploadBusy || S.source?.foreground_required && !S.source.foreground_ready) reason = ''; else if (!S.source) reason = 'Choose an example or upload an image.'; - $('st-generate').disabled = !signedIn || loading || failed || !S.source || S.uploadBusy || S.submitBusy || S.authBusy || + $('st-generate').disabled = inputBusy || !!geometry || !signedIn || loading || failed || !S.source || S.uploadBusy || S.authBusy || !!(S.source?.foreground_required && !S.source.foreground_ready); for (const id of ['st-file-input', 'st-pick-image', 'st-change-image', 'st-clear-image', 'st-seed']) - $(id).disabled = !signedIn || loading || failed; + $(id).disabled = inputBusy || !signedIn || loading || failed; $('st-history-toggle').disabled = !signedIn; $('st-history-more').disabled = !signedIn || S.pageLoading; - $('st-upload-retry').disabled = !signedIn || loading || failed || S.uploadBusy; - $('st-examples').inert = !signedIn; + $('st-upload-retry').disabled = inputBusy || !signedIn || loading || failed || S.uploadBusy; + $('st-examples').inert = !signedIn || inputBusy; $('st-examples').querySelectorAll('button[data-example]').forEach(gateCard); $('st-history-list').inert = !signedIn; $('st-history-list').querySelectorAll('button[data-job-id]').forEach(gateCard); if (!loading || failed || !signedIn) hideCardTip(); $('st-output-panel').inert = !signedIn || loading || failed; show('st-cancel-signin', S.authBusy); - text('st-generate', S.submitBusy ? 'Submitting…' : 'Generate 3D'); - const explanatory = loading || failed; + text('st-generate', S.submitBusy ? 'Submitting…' : geometry ? + texturesReady(geometry) ? 'Geometry generated' : 'Geometry ready' : 'Generate geometry'); + const explanatory = (loading || failed) && !imageFlow; show('st-generate-hit', explanatory); $('st-generate-hit').setAttribute('aria-label', reason); $('st-generate-hit').title = reason; - text('st-gate-note', reason); show('st-gate-note', !!reason); - show('st-native-startup', loading); - if (loading) { - const progress = workspaceLoading ? {} : S.deployment?.models?.startup_progress || {}; - const completed = Math.max(0, Number(progress.completed) || 0); - const total = Math.max(0, Number(progress.total) || 0); - const reconnecting = /^Reconnecting to model storage/.test(progress.label || ''); - $('st-native-startup').querySelector('strong').textContent = failed ? - workspaceLoading ? 'Workspace preparation failed' : 'Model loading failed' : - workspaceLoading ? 'Preparing your workspace…' : reconnecting ? 'Reconnecting to model storage…' : 'Loading models in the background'; - text('st-startup-detail', failed ? 'Please try again later' : total ? - progress.unit === 'bytes' ? `${(completed / 1073741824).toFixed(1)} / ${(total / 1073741824).toFixed(1)} GB` : - `${Math.min(completed, total)} / ${total} ${progress.unit === 'models' ? 'models' : 'stages'}` : 'Preparing…'); - text('st-startup-label', failed ? 'Service temporarily unavailable.' : 'Only required for the initial load.'); - show('st-startup-meter', !failed); - if (total) $('st-startup-meter').value = Math.min(100, completed / total * 100); - else $('st-startup-meter').removeAttribute('value'); - } + text('st-gate-note', reason); show('st-gate-note', !!reason && !imageFlow); requestAnimationFrame(() => { $('st-generate-hit').style.height = $('st-generate').offsetHeight + 'px'; }); } async function refreshStatus() { @@ -499,12 +593,17 @@ S.statusPromise = (async () => { try { const wasReady = ready(S.deployment); S.deployment = await api('/status'); updateGate(); + gpuDispatcher.syncRuntime(S.deployment.runtime_id); + gpuDispatcher.advance(record => preparationError(S.deployment, record)); + renderSelected(); + if (S.uploadState && currentUpload(S.uploadState) && S.uploadState.uploadId) void refreshUpload(S.uploadState); if (!wasReady && ready(S.deployment)) void activateWorkspace(); + return S.deployment; } catch (error) { if (!S.deployment) text('st-gate-note', 'Connecting to the service…'); } finally { S.statusPromise = null; clearTimeout(S.startupTimer); - if (!S.stopped && !ready(S.deployment) && !terminal(S.deployment)) + if (!S.stopped && pollStartup()) S.startupTimer = setTimeout(refreshStatus, document.hidden ? 20000 : 3000); } })(); @@ -516,20 +615,25 @@ } function historyStatus(job) { if (job.status !== 'complete') return {status: statusNames[job.status] ? job.status : 'running', - label: statusNames[job.status] || 'Preparing'}; + label: generationStatusName(job) || 'Preparing'}; + if (stagedJob(job) && !texturesReady(job)) { + const labels = {queued: 'Preparing textures', running: 'Generating textures', failed: 'Texture failed', interrupted: 'Texture interrupted'}; + return {status: labels[job.texture_status] ? job.texture_status : 'generation_complete', + label: labels[job.texture_status] || 'Geometry ready'}; + } // Generation completion makes the preview available. Export and recovery // remain separate operations, and only a downloadable GLB completes the result. const operations = [ - {status: job.export_status, labels: {queued: 'GLB Queued', running: 'Extracting GLB', + {status: job.export_status, labels: {queued: zeroGPURecord(job) ? 'Preparing GLB' : 'GLB Queued', running: 'Extracting GLB', failed: 'GLB Extraction Failed', interrupted: 'GLB Extraction Interrupted'}}, - {status: job.recover_status, labels: {queued: 'Preview Recovery Queued', running: 'Recovering Preview', + {status: job.recover_status, labels: {queued: zeroGPURecord(job) ? 'Preparing Preview' : 'Preview Recovery Queued', running: 'Recovering Preview', failed: 'Preview Recovery Failed', interrupted: 'Preview Recovery Interrupted'}} ]; const operation = operations.find(item => item.status === 'running') || operations.find(item => item.status === 'queued') || operations.find(item => ['failed', 'interrupted'].includes(item.status)); if (operation) return {status: operation.status, label: operation.labels[operation.status]}; - return privateURL(job.download_url) ? {status: 'complete', label: 'Complete'} : + return stagedJob(job) || privateURL(job.download_url) ? {status: 'complete', label: 'Complete'} : {status: 'generation_complete', label: 'Generation Complete'}; } function renderHistory() { @@ -591,7 +695,7 @@ ++S.selectionVersion; S.previewController?.abort(); S.previewController = null; S.selected = null; S.previewLoaded = null; clearPreview(); // A newer upload can coexist with the previously selected result. - if (job?.upload_id && S.source?.upload_id === job.upload_id) { + if (!leftInputBusy() && job?.upload_id && S.source?.upload_id === job.upload_id) { ++S.inputVersion; clearUploadPreview(); S.source = null; S.submitKey = null; if (S.localURL) URL.revokeObjectURL(S.localURL); S.localURL = null; $('st-input-image').removeAttribute('src'); $('st-file-input').value = ''; @@ -675,39 +779,59 @@ S.historyError = sustained ? (S.historyLoaded ? 'History could not refresh. We will keep trying.' : 'History could not load. We will keep trying.') : ''; } - } finally { S.refreshPromise = null; S.pageLoading = false; renderHistory(); scheduleJobs(); } + } finally { S.refreshPromise = null; S.pageLoading = false; renderHistory(); updateGate(); scheduleJobs(); } })(); return S.refreshPromise; } function renderSelected() { - if (!ready(S.deployment)) return; + updateGate(); if (S.selected?.kind !== 'job') return; const job = S.jobs.get(S.selected.id); if (!job) return; + if (!ready(S.deployment)) { + if (job.status === 'complete') renderTexture(job); + return; + } if (job.status === 'complete') { - if (S.previewLoaded !== previewKey(job)) void loadPreview(job); - renderExport(job); return; + if (!textureProgressVisible(job) && S.previewLoaded !== previewKey(job)) void loadPreview(job); + show('st-job-progress', false); renderTexture(job); renderExport(job); return; } + $('st-output-panel').dataset.textureProgress = 'false'; + show('st-texture-actions', false); show('st-extract', !stagedJob(job)); show('st-export-note', !stagedJob(job)); S.previewLoaded = null; $('st-model').replaceChildren(); $('st-ovoxel').replaceChildren(); show('st-model', false); show('st-ovoxel', false); show('st-native-empty', false); show('st-job-progress', true); $('st-extract').disabled = true; show('st-download', false); show('st-export-progress', false); clearTimeout(S.exportTimer); S.exportTimer = null; const progress = job.progress || {}; - text('st-job-title', statusNames[job.status] || 'Preparing your model'); + text('st-job-title', generationStatusName(job) || 'Preparing your model'); const position = progress.queue_position; - const detail = job.status === 'failed' || job.status === 'interrupted' ? - (job.status === 'interrupted' ? 'This generation was interrupted. You can generate it again.' : 'This generation could not finish. You can try again.') : + const gpuFailure = gpuDispatcher.failure(job.gpu_operation_id); + const resumeNeeded = needsGPUResume(job); + const resuming = !!S.resumeOperation && S.resumeOperation === job.gpu_operation_id; + $('st-job-progress').dataset.resumable = String(!!resumeNeeded); + show('st-resume-gpu', !!resumeNeeded); + resumeGPU.disabled = resuming; + text('st-resume-gpu', resuming ? 'Resuming…' : gpuFailure ? 'Retry task' : 'Resume task'); + const detail = resuming ? 'Reconnecting to your task…' : gpuFailure ? gpuFailure : resumeNeeded ? 'Your task is saved. Resume when you are ready.' : + job.status === 'failed' || job.status === 'interrupted' ? + (job.status === 'interrupted' ? 'This generation was interrupted. You can generate it again.' : + englishMessage(job.error, 'This generation could not finish. You can try again.')) : + taskPreparing(job) ? preparationText(job) : + job.status === 'queued' && zeroGPURecord(job) ? (stagedJob(job) ? 'Starting geometry… Keep this page open.' : 'Starting generation… Keep this page open.') : job.status === 'queued' && position !== undefined && position !== null ? - `Waiting in queue · Position ${position}` : englishMessage(progress.label, 'Preparing your model…'); + `Waiting in queue · Position ${position}` : generationProgressLabel(progress.label, 'Preparing your model…'); text('st-job-detail', detail); - show('st-queue-note', job.status === 'queued'); - $('st-job-progress').querySelector('progress').hidden = !['queued', 'running'].includes(job.status); - text('st-job-time', progress.elapsed_seconds ? `${Math.floor(progress.elapsed_seconds)}s elapsed` : ''); + show('st-queue-note', job.status === 'queued' && !zeroGPURecord(job)); + $('st-job-progress').querySelector('progress').hidden = !!resumeNeeded || !['queued', 'running'].includes(job.status); + preparationMeter($('st-job-progress').querySelector('progress'), job); + const timing = !resumeNeeded && (taskPreparing(job) ? 'Only needed for the first generation after each Space startup.' : + progress.elapsed_seconds ? `${Math.floor(progress.elapsed_seconds)}s elapsed` : ''); + text('st-job-time', timing); show('st-job-time', !!timing); text('st-export-note', ''); } function previewKey(job) { - const view = privateURL(job.download_url) ? S.selected?.view || 'glb' : 'render'; - return `${view}:${job.id}`; + const view = texturesReady(job) && privateURL(job.download_url) ? S.selected?.view || 'glb' : 'render'; + return `${view}:${job.id}:${job.asset_id || ''}`; } async function loadPreview(job) { const version = S.selectionVersion, inputVersion = S.inputVersion; @@ -720,52 +844,141 @@ const view = key.split(':')[0]; const result = await api('/jobs/' + encodeURIComponent(job.id) + '/preview?view=' + view, {signal: controller.signal}); if (version !== S.selectionVersion || S.selected?.id !== job.id || S.previewController !== controller) return; + const current = S.jobs.get(job.id) || job; + if (previewKey(current) !== key) { S.previewLoaded = null; renderSelected(); return; } + const textureProgress = textureProgressVisible(current); $('st-ovoxel').innerHTML = result.html || ''; - show('st-ovoxel', !!result.html); show('st-native-empty', !result.html); + show('st-ovoxel', !!result.html && !textureProgress); show('st-native-empty', !result.html && !textureProgress); if (!result.html) text('st-native-empty', 'The saved preview is not available.'); const input = privateURL(result.display_url || job.display_url || result.input_url || job.input_url); const original = privateURL(result.input_url || job.input_url); - if (input && !S.source?.local_file && !S.source?.local_job) + if (input && (!leftInputBusy() || !S.source) && !S.source?.local_file && !S.source?.local_job) setInput(input, {upload_id: result.upload_id || job.upload_id, input_url: original}, result.seed ?? job.seed, inputVersion); - const current = S.jobs.get(job.id) || job; if (result.download_url) current.download_url = result.download_url; const model = $('st-ovoxel').querySelector('.st-gallery-preview[data-model]'); - S.previewLoaded = result.view ? `${result.view}:${job.id}` : + S.previewLoaded = result.view ? `${result.view}:${job.id}:${job.asset_id || ''}` : model && privateURL(model.dataset.model) === privateURL(current.download_url) ? - `glb:${job.id}` : `render:${job.id}`; - renderHistory(); renderExport(current); + `glb:${job.id}:${job.asset_id || ''}` : `render:${job.id}:${job.asset_id || ''}`; + renderHistory(); renderTexture(current); renderExport(current); } catch (error) { if (quietError(error) || version !== S.selectionVersion || S.previewController !== controller) return; - S.previewLoaded = null; text('st-native-empty', 'Could not load this preview. Select it again to retry.'); + S.previewLoaded = null; + if (textureProgressVisible(S.jobs.get(job.id) || job)) return; + text('st-native-empty', 'Could not load this preview. Select it again to retry.'); notice(errorText(error)); } } + function renderTexture(job) { + const visible = stagedJob(job) && job.status === 'complete' && !texturesReady(job); + const central = visible && textureProgressVisible(job); + $('st-output-panel').dataset.textureProgress = String(!!central); + show('st-texture-actions', visible); + if (!visible) return; + const previewReady = S.previewLoaded === previewKey(job) && !!$('st-ovoxel').innerHTML; + show('st-model', false); show('st-ovoxel', !central && previewReady); + show('st-native-empty', !central && !previewReady); show('st-job-progress', false); + const active = ['queued', 'running'].includes(job.texture_status); + const busy = S.textureBusy === job.id; + const gpuFailure = textureFailure(job); + const resumeNeeded = job.texture_status === 'queued' && job.gpu_operation_kind === 'texture' && needsGPUResume(job); + const failed = ['failed', 'interrupted'].includes(job.texture_status); + const preparing = active && taskPreparing(job); + const progress = job.texture_progress || {}; + show('st-texture-note', !active && !busy && !failed && !gpuFailure); + show('st-texture-progress', active || busy || failed || !!gpuFailure); + text('st-texture-title', busy ? 'Preparing textures' : failed || gpuFailure ? 'Texture failed' : + job.texture_status === 'running' ? 'Generating textures' : 'Preparing textures'); + text('st-texture-detail', busy ? 'Connecting to your task…' : gpuFailure || (resumeNeeded ? + 'Your geometry is ready. Resume textures when you are ready.' : failed ? + englishMessage(job.texture_error, 'Texture generation could not finish. Retry textures to continue.') : preparing ? + preparationText(job) : job.texture_status === 'running' ? generationProgressLabel(progress.label, 'Generating textures…') : + 'Starting textures… Keep this page open.')); + const meter = $('st-texture-progress').querySelector('progress'); + meter.hidden = !busy && (!active || !!resumeNeeded || !!gpuFailure); + preparationMeter(meter, job); + const timing = !preparing && !resumeNeeded && !busy && progress.elapsed_seconds ? `${Math.floor(progress.elapsed_seconds)}s elapsed` : ''; + text('st-texture-time', timing); show('st-texture-time', !!timing); + show('st-texture', !active || !!resumeNeeded || busy); + $('st-texture').disabled = busy || active && !resumeNeeded || !S.identity?.authenticated || !ready(S.deployment); + text('st-texture', busy ? 'Connecting…' : gpuFailure || failed ? 'Retry textures' : resumeNeeded ? 'Resume textures' : 'Generate textures'); + } + async function startTexture() { + const job = S.selected?.kind === 'job' && S.jobs.get(S.selected.id); + if (!job || !stagedJob(job) || job.status !== 'complete' || texturesReady(job) || + !S.identity?.authenticated || S.textureBusy || $('st-texture').disabled) return; + const identityVersion = S.identityVersion; + S.textureErrors.delete(job.id); + S.textureBusy = job.id; updateGate(); renderTexture(job); + try { + const deployment = await refreshStatus(); + if (identityVersion !== S.identityVersion || S.stopped || !S.identity?.authenticated) return; + if (!deployment || !ready(deployment)) throw new Error('The service is reconnecting. Please try again.'); + const current = S.jobs.get(job.id); + if (!current || texturesReady(current)) return; + if (current.texture_status === 'queued') { + if (current.gpu_operation_kind === 'texture' && needsGPUResume(current)) void gpuDispatcher.retry(current); + return; + } + if (current.texture_status === 'running') return; + if (!S.textureKeys.has(job.id)) S.textureKeys.set(job.id, crypto.randomUUID()); + const updated = await api('/jobs/' + encodeURIComponent(job.id) + '/texture', + {method: 'POST', json: {idempotency_key: S.textureKeys.get(job.id)}}); + if (identityVersion !== S.identityVersion || S.stopped) return; + S.textureKeys.delete(job.id); S.jobs.set(updated.id, updated); + renderHistory(); scheduleJobs(); + } catch (error) { + if (identityVersion === S.identityVersion && !quietError(error)) { + const message = errorText(error); S.textureErrors.set(job.id, message); + if (S.selected?.kind !== 'job' || S.selected.id !== job.id) notice(message); + } + } finally { + if (identityVersion === S.identityVersion) { + S.textureBusy = null; + updateGate(); + if (S.selected?.kind === 'job' && S.selected.id === job.id) renderSelected(); + } + } + } function renderExport(job) { clearTimeout(S.exportTimer); S.exportTimer = null; + const allowed = texturesReady(job); + show('st-extract', allowed); show('st-export-note', allowed); + if (!allowed) { + $('st-extract').disabled = true; show('st-download', false); show('st-export-progress', false); + show('st-view-switch', false); text('st-export-note', ''); return; + } const exporting = ['queued', 'running'].includes(job.export_status); + const gpuFailure = job.export_status === 'queued' && gpuDispatcher.failure(job.gpu_operation_id); + const resumeNeeded = job.export_status === 'queued' && needsGPUResume(job); const url = privateURL(job.download_url); show('st-view-switch', !!url && job.status === 'complete'); const view = previewKey(job).split(':')[0]; $('st-view-switch').querySelectorAll('[data-result-view]').forEach(button => button.setAttribute('aria-pressed', String(button.dataset.resultView === view))); - $('st-extract').disabled = (!url && exporting) || job.status !== 'complete' || !S.identity?.authenticated || !ready(S.deployment); + $('st-extract').disabled = (!url && exporting && !resumeNeeded) || job.status !== 'complete' || !S.identity?.authenticated || !ready(S.deployment); $('st-extract').dataset.action = url ? 'download' : 'extract'; - text('st-extract', url ? 'Download GLB' : exporting ? - job.export_status === 'queued' ? 'GLB extraction queued…' : 'Extracting GLB…' : 'Extract GLB'); + text('st-extract', url ? 'Download GLB' : gpuFailure ? 'Retry extraction' : resumeNeeded ? 'Resume GLB extraction' : exporting ? + job.export_status === 'queued' ? zeroGPURecord(job) ? 'Preparing GLB…' : 'GLB extraction queued…' : 'Extracting GLB…' : 'Extract GLB'); show('st-download', false); if (url) $('st-download').href = url; else $('st-download').removeAttribute('href'); show('st-export-progress', exporting && !url); - text('st-export-note', url ? 'Your GLB is ready to download.' : exporting ? + text('st-export-note', url ? 'Your GLB is ready to download.' : gpuFailure ? gpuFailure : resumeNeeded ? + 'Your extraction request is saved. Resume it when you are ready.' : taskPreparing(job) ? preparationText(job) : exporting ? 'You can return to My History while extraction continues.' : - job.export_status === 'failed' ? 'Extraction failed. Click Extract GLB to try again.' : 'Extract and download GLB.'); + job.export_status === 'failed' ? englishMessage(job.export_error, 'Extraction failed. Click Extract GLB to try again.') : 'Extract and download GLB.'); if (!exporting || url) return; - const progress = job.export_progress || {}, queued = job.export_status === 'queued'; + const preparing = taskPreparing(job), queued = job.export_status === 'queued'; + const progress = (preparing ? preparationState(S.deployment, job).startup_progress : job.export_progress) || {}; const position = progress.queue_position; - text('st-export-label', queued ? `Waiting for GPU${position ? ' · Queue position ' + position : ''}` : + text('st-export-label', preparing ? preparationText(job) : resumeNeeded ? 'Extraction paused' : + queued ? zeroGPURecord(job) ? 'Preparing GLB… Keep this page open.' : `Waiting for extraction${position ? ' · Queue position ' + position : ''}` : englishMessage(progress.label, 'Preparing your GLB…')); const completed = Math.max(0, Number(progress.completed) || 0), total = Math.max(0, Number(progress.total) || 0); const meter = $('st-export-meter'), details = []; - if (total && !queued) { + if (preparing) { + meter.value = formatProgress(progress).percent; + meter.setAttribute('aria-valuetext', formatProgress(progress).percent + '%'); + } else if (total && !queued) { meter.value = Math.min(100, completed / total * 100); meter.setAttribute('aria-valuetext', `${Math.min(completed, total)} of ${total} ${progress.unit === 'bytes' ? 'bytes' : 'stages'} complete`); details.push(progress.unit === 'bytes' ? @@ -773,13 +986,13 @@ `${Math.min(completed, total)} / ${total} stages`); } else { meter.removeAttribute('value'); meter.removeAttribute('aria-valuetext'); } const transferred = Number(progress.bytes_completed), transferTotal = Number(progress.bytes_total); - if (transferTotal > 0 && progress.unit !== 'bytes') details.push( + if (!preparing && transferTotal > 0 && progress.unit !== 'bytes') details.push( `${(Math.max(0, transferred || 0) / 1048576).toFixed(1)} / ${(transferTotal / 1048576).toFixed(1)} MB transferred`); job._exportProgressSeenAt ||= Date.now(); const elapsed = Math.max(0, Number(progress.elapsed_seconds) || 0) + (progress.clock === 'running' ? Math.floor((Date.now() - job._exportProgressSeenAt) / 1000) : 0); if (elapsed) details.push(`${Math.floor(elapsed)}s elapsed`); - details.push(queued ? 'Extraction shares the generation queue.' : 'Stage durations vary with model complexity.'); + if (!preparing && queued && !zeroGPURecord(job)) details.push('Extraction shares the generation queue.'); text('st-export-detail', details.join(' · ')); S.exportTimer = setTimeout(() => { if (S.selected?.kind === 'job' && S.selected.id === job.id) renderExport(S.jobs.get(job.id) || job); @@ -788,23 +1001,25 @@ function selectJob(id, keepInput = false) { if (!S.identity?.authenticated || !ready(S.deployment)) return; const job = S.jobs.get(id); if (!job) return; - ++S.selectionVersion; ++S.inputVersion; S.previewController?.abort(); + const preserveReference = leftInputBusy() && !!S.source; + ++S.selectionVersion; if (!preserveReference) ++S.inputVersion; + S.previewController?.abort(); S.selected = {kind: 'job', id}; S.previewLoaded = null; clearPreview(); const input = privateURL(job.display_url || job.input_url) || `/api/history/jobs/${encodeURIComponent(id)}/files/input.png`; const original = privateURL(job.input_url) || `/api/history/jobs/${encodeURIComponent(id)}/files/input.png`; const currentInput = keepInput && $('st-input-image').getAttribute('src'); - if (currentInput && S.source) { + if (!preserveReference && currentInput && S.source) { // A submission already has a visible reference, including examples. // Keep it while the newly saved private files become available. const source = {...S.source}; delete source.example; setInput(currentInput, {...source, local_job: id, upload_id: job.upload_id, input_url: original}, job.seed); - } else setInput(input, {upload_id: job.upload_id, input_url: original}, job.seed); + } else if (!preserveReference) setInput(input, {upload_id: job.upload_id, input_url: original}, job.seed); renderHistory(); renderSelected(); if (S.expanded && matchMedia('(max-width:900px)').matches) expandHistory(false); } function selectExample(name) { - if (!S.identity?.authenticated || !ready(S.deployment)) return; + if (leftInputBusy() || !S.identity?.authenticated || !ready(S.deployment)) return; const example = S.examples.get(name); if (!example) return; ++S.selectionVersion; ++S.inputVersion; S.previewController?.abort(); S.selected = {kind: 'example', id: name}; S.previewLoaded = null; @@ -863,21 +1078,20 @@ $('st-upload-status').dataset.state = phase; text('st-upload-label', label); show('st-upload-status', true); show('st-upload-retry', retry && (!!state.uploadId || state.retryAction === 'upload')); - show('st-upload-meter', !['complete', 'failed'].includes(phase)); - $('st-upload-meter').removeAttribute('value'); updateGate(); + updateGate(); } function scheduleUploadStatus(state, delay = document.hidden ? 6000 : 1500) { clearTimeout(state.timer); if (currentUpload(state)) state.timer = setTimeout(() => void refreshUpload(state), delay); } function uploadReadFailed(state, error) { - if (!currentUpload(state) || quietError(error)) return; + if (!currentUpload(state) || state.appliedResultKey || state.applyPromise || quietError(error)) return; if (state.resultURL) { URL.revokeObjectURL(state.resultURL); state.resultURL = null; } state.failures++; state.failureSince ??= Date.now(); if (state.failures < 3 || Date.now() - state.failureSince < 30000) { S.uploadBusy = true; - uploadLabel(state, 'checking', 'Checking background-removal progress…'); + uploadLabel(state, 'checking', 'Preparing image…'); scheduleUploadStatus(state, Math.min(8000, state.failures * 2000)); } else { S.uploadBusy = false; state.retryAction = 'status'; @@ -888,9 +1102,20 @@ if (!currentUpload(state)) return; if (!result.upload_id || state.uploadId && state.uploadId !== result.upload_id) throw new Error('The uploaded image could not be verified.'); + // One upload state belongs to one account and input version. Repeated + // status/dispatcher notifications must not re-download its applied cutout. + const foreground = result.foreground; + const resultKey = foreground?.status === 'complete' ? JSON.stringify([ + result.upload_id, result.files?.['foreground.png']?.version || result.foreground_url + ]) : null; + if (state.appliedResultKey && (!resultKey || state.appliedResultKey === resultKey)) return; + if (state.applyPromise) { + await state.applyPromise; + return applyUploadResult(result, state); + } state.uploadId = result.upload_id; + state.gpuRecord = result; S.source = {...S.source, upload_id: result.upload_id, uploaded: true}; - const foreground = result.foreground; if (!foreground) { // Retain compatibility with saved uploads and the previous API version. S.source.foreground_required = false; S.uploadBusy = false; @@ -898,53 +1123,74 @@ } S.source.foreground_required = true; S.source.foreground_ready = false; if (['queued', 'running'].includes(foreground.status)) { + const gpuFailure = gpuDispatcher.failure(result.gpu_operation_id); + if (gpuFailure) { + S.uploadBusy = false; state.retryAction = 'gpu'; state.gpuRecord = result; + clearTimeout(state.timer); uploadLabel(state, 'failed', gpuFailure, true); return; + } state.failures = 0; state.failureSince = null; S.uploadBusy = true; - const position = Number(foreground.progress?.queue_position); - uploadLabel(state, foreground.status, foreground.status === 'queued' ? - 'Waiting to remove background…' + (position > 0 ? ' · Queue position ' + position : '') : 'Removing background…'); + uploadLabel(state, foreground.status, taskPreparing(result) ? 'Loading model…' : foreground.status === 'queued' ? + 'Removing background… Keep this page open.' : 'Removing background…'); scheduleUploadStatus(state); return; } if (['failed', 'interrupted'].includes(foreground.status)) { S.uploadBusy = false; state.retryAction = 'foreground'; state.retryKey = null; - uploadLabel(state, 'failed', 'Background removal failed. Please retry.', true); return; + uploadLabel(state, 'failed', englishMessage(foreground.error, 'Background removal failed. Please retry.'), true); return; } if (foreground.status !== 'complete') throw new Error('Background removal status is unavailable.'); - S.uploadBusy = true; uploadLabel(state, 'loading', 'Loading background-removed image…'); - const url = privateURL(result.foreground_url); - const expected = '/api/history/uploads/' + encodeURIComponent(state.uploadId) + '/files/foreground.png'; - if (!url || new URL(url).pathname !== expected) throw new Error('The background-removed image is unavailable.'); - const response = await fetch(url, {credentials: 'same-origin', cache: 'no-store', - signal: AbortSignal.any([state.controller.signal, AbortSignal.timeout(30000)])}); - if (!response.ok) throw new Error('The background-removed image could not load.'); - const blob = await response.blob(); - if (!currentUpload(state)) return; - if (!blob.size || blob.size > 8 * 1024 * 1024 || !blob.type.startsWith('image/')) - throw new Error('The background-removed image is invalid.'); - state.resultURL = URL.createObjectURL(blob); - const image = new Image(); image.src = state.resultURL; - await image.decode(); - if (!currentUpload(state)) return; - const previous = S.localURL; - S.localURL = state.resultURL; state.resultURL = null; - if (previous) URL.revokeObjectURL(previous); - S.uploadBusy = false; state.failures = 0; state.failureSince = null; state.retryKey = null; - setInput(S.localURL, {...S.source, foreground_required: true, foreground_ready: true, - foreground_file: blob}, undefined, state.inputVersion); - uploadLabel(state, 'complete', 'Background removed.'); + const applying = (async () => { + S.uploadBusy = true; uploadLabel(state, 'loading', 'Loading image…'); + const url = privateURL(result.foreground_url); + const expected = '/api/history/uploads/' + encodeURIComponent(state.uploadId) + '/files/foreground.png'; + if (!url || new URL(url).pathname !== expected) throw new Error('The background-removed image is unavailable.'); + const response = await fetch(url, {credentials: 'same-origin', cache: 'no-store', + signal: AbortSignal.any([state.controller.signal, AbortSignal.timeout(30000)])}); + if (!response.ok) throw new Error('The background-removed image could not load.'); + const blob = await response.blob(); + if (!currentUpload(state)) return; + if (!blob.size || blob.size > 8 * 1024 * 1024 || !blob.type.startsWith('image/')) + throw new Error('The background-removed image is invalid.'); + state.resultURL = URL.createObjectURL(blob); + const image = new Image(); image.src = state.resultURL; + await image.decode(); + if (!currentUpload(state)) return; + const previous = S.localURL; + S.localURL = state.resultURL; state.resultURL = null; + if (previous) URL.revokeObjectURL(previous); + S.uploadBusy = false; state.failures = 0; state.failureSince = null; state.retryKey = null; + state.appliedResultKey = resultKey; clearTimeout(state.timer); state.timer = null; + setInput(S.localURL, {...S.source, foreground_required: true, foreground_ready: true, + foreground_file: blob}, undefined, state.inputVersion); + uploadLabel(state, 'complete', 'Background removed.'); + })(); + state.applyPromise = applying; + try { await applying; } + finally { if (state.applyPromise === applying) state.applyPromise = null; } } async function refreshUpload(state) { - if (!currentUpload(state) || !state.uploadId) return; - try { - const result = await api('/uploads/' + encodeURIComponent(state.uploadId), {signal: state.controller.signal}); - await applyUploadResult(result, state); - } catch (error) { uploadReadFailed(state, error); } + if (!currentUpload(state) || !state.uploadId || state.appliedResultKey) return; + if (state.refreshPromise) return state.refreshPromise; + state.refreshPromise = (async () => { + try { + const result = await api('/uploads/' + encodeURIComponent(state.uploadId), {signal: state.controller.signal}); + await applyUploadResult(result, state); + } catch (error) { uploadReadFailed(state, error); } + finally { state.refreshPromise = null; } + })(); + return state.refreshPromise; } async function retryUploadPreview() { const state = S.uploadState; - if (!state || !currentUpload(state) || S.uploadBusy || !ready(S.deployment)) return; + if (leftInputBusy() || !state || !currentUpload(state) || S.uploadBusy || !ready(S.deployment)) return; if (state.retryAction === 'upload') return upload(S.source?.local_file); + state.appliedResultKey = null; + if (state.retryAction === 'gpu' && state.gpuRecord) { + S.uploadBusy = true; + uploadLabel(state, 'queued', 'Preparing image…'); + void gpuDispatcher.retry(state.gpuRecord); scheduleUploadStatus(state); return; + } state.failures = 0; state.failureSince = null; S.uploadBusy = true; - uploadLabel(state, 'queued', 'Waiting to remove background…'); + uploadLabel(state, 'queued', 'Preparing image…'); if (state.retryAction !== 'foreground') return refreshUpload(state); state.retryKey ||= crypto.randomUUID(); try { @@ -954,7 +1200,7 @@ } catch (error) { uploadReadFailed(state, error); } } async function upload(file) { - if (!file) return; + if (!file || leftInputBusy()) return; if (!S.identity?.authenticated || !ready(S.deployment) || S.stopped) return; if (file.size > 20 * 1024 * 1024) return notice('Please upload an image smaller than 20 MiB.'); const serviceWaits = S.source?.local_file === file && S.uploadState?.phase === 'waiting_service' ? @@ -968,7 +1214,7 @@ S.localURL = URL.createObjectURL(file); setInput(S.localURL, {local_file: file, foreground_required: true, foreground_ready: false}, undefined, version); S.uploadBusy = true; S.submitKey = null; updateGate(); - uploadLabel(state, 'uploading', 'Uploading your image…'); + uploadLabel(state, 'uploading', 'Loading image…'); try { const form = new FormData(); form.append('file', file); const result = await api('/upload', {method: 'POST', body: form, signal: state.controller.signal}); @@ -982,20 +1228,20 @@ S.uploadBusy = false; state.serviceWaits++; state.retryAction = 'upload'; S.deployment = {...S.deployment, workspace: {stage: 'starting'}}; uploadLabel(state, state.serviceWaits <= 3 ? 'waiting_service' : 'failed', - state.serviceWaits <= 3 ? 'Waiting for the service. Your image is kept here.' : + state.serviceWaits <= 3 ? 'Preparing image…' : 'The service is still preparing. Your image is kept here. Please retry when ready.', state.serviceWaits > 3); void refreshStatus(); } else { S.uploadBusy = false; state.retryAction = 'upload'; - uploadLabel(state, 'failed', 'Upload could not finish. Your image is kept here. Please retry.', true); - notice(errorText(error)); + uploadLabel(state, 'failed', errorText(error), true); } } } finally { if (currentUpload(state)) { $('st-file-input').value = ''; updateGate(); } } } function stageLocalFile(file, seed) { + if (leftInputBusy()) return; clearUploadPreview(); ++S.inputVersion; S.uploadBusy = false; if (S.localURL) URL.revokeObjectURL(S.localURL); @@ -1005,7 +1251,7 @@ if (S.identity?.authenticated && ready(S.deployment)) void upload(file); } async function submit() { - if ($('st-generate').disabled || S.submitBusy) return; + if (leftInputBusy() || selectedGeometry() || $('st-generate').disabled) return; if (!S.identity?.authenticated) return connectAccount(); const source = {...S.source}, seed = Math.round(Number($('st-seed').value)); if (!Number.isFinite(seed)) return notice('Please enter a valid seed.'); @@ -1060,6 +1306,10 @@ async function extract() { if (S.selected?.kind !== 'job' || $('st-extract').disabled) return; const job = S.jobs.get(S.selected.id); if (!job) return; + if (!texturesReady(job)) return; + if (job.export_status === 'queued' && needsGPUResume(job)) { + void gpuDispatcher.retry(job); renderExport(job); return; + } const url = privateURL(job.download_url); if (url) { try { @@ -1104,20 +1354,25 @@ root.querySelector('#st-generate-control').dataset.notice = 'true'; }); $('st-generate-hit').addEventListener('mouseleave', () => delete root.querySelector('#st-generate-control').dataset.notice); - const pick = () => { if (ready(S.deployment) && S.identity?.authenticated) $('st-file-input').click(); }; + const pick = () => { if (!leftInputBusy() && ready(S.deployment) && S.identity?.authenticated) $('st-file-input').click(); }; $('st-pick-image').addEventListener('click', pick); $('st-change-image').addEventListener('click', pick); $('st-file-input').addEventListener('change', event => void upload(event.target.files[0])); $('st-upload').addEventListener('dragover', event => event.preventDefault()); $('st-upload').addEventListener('drop', event => { event.preventDefault(); void upload(event.dataTransfer.files[0]); }); + for (const eventName of ['dragover', 'drop']) document.addEventListener(eventName, event => { + if (leftInputBusy() && Array.from(event.dataTransfer?.types || []).includes('Files')) event.preventDefault(); + }, true); $('st-clear-image').addEventListener('click', () => { - if (!ready(S.deployment)) return; + if (leftInputBusy() || !ready(S.deployment)) return; clearUploadPreview(); ++S.inputVersion; S.source = null; S.uploadBusy = false; S.submitKey = null; if (S.localURL) URL.revokeObjectURL(S.localURL); S.localURL = null; $('st-input-image').removeAttribute('src'); show('st-input-image', false); show('st-upload-placeholder', true); show('st-clear-image', false); updateGate(); }); $('st-upload-retry').addEventListener('click', () => void retryUploadPreview()); - $('st-seed').addEventListener('input', () => { S.submitKey = null; }); + $('st-seed').addEventListener('input', () => { + if (!leftInputBusy()) { S.submitKey = null; updateGate(); } + }); $('st-generate').addEventListener('click', () => void submit()); const cancelSignIn = () => accountConnectionFailed('Sign-in cancelled. Please sign in to continue.'); $('st-cancel-signin').addEventListener('click', cancelSignIn); @@ -1125,13 +1380,14 @@ $('st-auth-signin').addEventListener('click', () => connectAccount(true)); $('st-auth-dialog').addEventListener('cancel', event => event.preventDefault()); $('st-extract').addEventListener('click', () => void extract()); + $('st-texture').addEventListener('click', () => startTexture()); $('st-view-switch').addEventListener('click', event => { const button = event.target.closest('[data-result-view]'); if (!button || S.selected?.kind !== 'job') return; const job = S.jobs.get(S.selected.id); if (!job || job.status !== 'complete' || !privateURL(job.download_url)) return; const view = button.dataset.resultView; - if (!['render', 'glb'].includes(view) || previewKey(job) === `${view}:${job.id}`) return; + if (!['render', 'glb'].includes(view) || previewKey(job) === `${view}:${job.id}:${job.asset_id || ''}`) return; S.selected.view = view; ++S.selectionVersion; S.previewLoaded = null; renderSelected(); @@ -1215,7 +1471,7 @@ } document.addEventListener('visibilitychange', () => { clearTimeout(S.jobsTimer); clearTimeout(S.startupTimer); - if (document.hidden) { scheduleJobs(); if (!ready(S.deployment) && !terminal(S.deployment)) S.startupTimer = setTimeout(refreshStatus, 20000); } + if (document.hidden) { scheduleJobs(); if (pollStartup()) S.startupTimer = setTimeout(refreshStatus, 20000); } else void resume(); }); window.addEventListener('focus', () => void resume()); diff --git a/hf_space/history_ui.py b/hf_space/history_ui.py index d3d7a719843aecdf604e8c32abb62db1c6abd042..89a4ea8048a79b3fab11104d1ec8367dff56442b 100644 --- a/hf_space/history_ui.py +++ b/hf_space/history_ui.py @@ -1,4 +1,5 @@ """Native, account-scoped UI. Private images never enter Gradio's file cache.""" +import base64 from pathlib import Path import gradio as gr from hf_space.gallery_preview import install_js as gallery_install_js @@ -9,9 +10,20 @@ CSS = (DIRECTORY / "ui.css").read_text(encoding="utf-8") + "\n" + ( DIRECTORY / "history_ui.css").read_text(encoding="utf-8") HEADER = ('

Supra-Tex

' '

3D Generation with Super-Aligned Texture

') -RESOURCE_NOTICE = ('') +# Official HF artwork and PRO badge treatment; FREE has no separate plan logo. +# Sources: https://huggingface.co/brand and https://huggingface.co/pricing +# Logo: https://huggingface.co/front/assets/huggingface_logo-noborder.svg +HF_LOGO = base64.b64encode((DIRECTORY / 'brand/huggingface-logo.svg').read_bytes()).decode('ascii') +RESOURCE_NOTICE = ( + '' +) # Keep Gradio's footer and settings controls in English on Chinese browsers. # Flat translation keys take precedence over Gradio's nested locale resources. ENGLISH_UI = { @@ -64,15 +76,14 @@ MARKUP = """
- +
Generation settings
- +

Checking service…

-
@@ -88,6 +99,11 @@ MARKUP = """ + @@ -134,7 +150,8 @@ def build_ui(deployment_status=None): gr.LoginButton(visible=False, elem_id="st-oauth-bootstrap") gr.HTML(MARKUP, apply_default_css=False, elem_id="st-native-page") gr.HTML(RESOURCE_NOTICE, apply_default_css=False, elem_id="st-resource-footer") - script = (DIRECTORY / "history_ui.js").read_text(encoding="utf-8") + script = (DIRECTORY / "zerogpu_client.js").read_text(encoding="utf-8") + "\n" + ( + DIRECTORY / "history_ui.js").read_text(encoding="utf-8") demo.load(None, js="() => {\n" + script + "\n}", api_visibility="private", queue=False) demo.load(None, js=gallery_install_js(), api_visibility="private", queue=False) for filename in ("ovoxel_preview.js", "viewport.js"): diff --git a/hf_space/native_build.py b/hf_space/native_build.py index c19b5b831c0f1d212203972b595af505ef8b3781..2681eed61f31dfbf57e94ca86f69d9c49577d059 100644 --- a/hf_space/native_build.py +++ b/hf_space/native_build.py @@ -7,24 +7,47 @@ import importlib.metadata import copy import json import os -from pathlib import Path +from pathlib import Path, PurePosixPath import shutil import site import subprocess import sys import threading +import tarfile import time import traceback import uuid import zipfile from scripts.install_runtime import O_VOXEL_FLAGS, SOURCES +from hf_space.gpu_mode import enabled as zero_gpu_enabled + + +def cuda_architectures(): + if not zero_gpu_enabled(): + return '8.9+PTX' + # The scheduler advertises the CUDA architecture without acquiring a GPU. + # Cover both shared fleets if an older image does not publish these values. + major = os.environ.get('ZEROGPU_CUDA_CAPABILITY_MAJOR') + minor = os.environ.get('ZEROGPU_CUDA_CAPABILITY_MINOR', '0') + if major is None: + return '9.0;12.0+PTX' + if major not in ('8', '9', '10', '12') or minor not in ('0', '6', '9'): + raise RuntimeError('Unsupported advertised ZeroGPU CUDA architecture') + return f'{major}.{minor}+PTX' + + +def wheel_cache_key(): + if not zero_gpu_enabled(): + return 'torch280-cu128-cp310-sm89-v1' + arch = cuda_architectures().replace('.', '').replace(';', '-').replace('+', '-') + return f'torch280-cu128-cp310-zero-sm{arch}-v1' ROOT = Path(__file__).resolve().parents[1] -WORK = Path('/tmp/supratex-build-cu128') +WORK = Path('/tmp/supratex-build-cu128-zero' if zero_gpu_enabled() else '/tmp/supratex-build-cu128') STATE = Path(os.environ.get('SUPRATEX_NATIVE_STATE', str(WORK / 'status.json'))) LOG = STATE.parent / 'build.log' -CACHE = Path('/data/native-wheels/torch280-cu128-cp310-sm89-v1') +CACHE = Path('/data/native-wheels') / wheel_cache_key() LOCK = threading.Lock() _persist_lock = threading.Lock() _start_lock = threading.Lock() @@ -45,6 +68,78 @@ CUDA_PARTS = ( ('cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.8.90-archive.tar.xz', '8d566b5fe745c46842dc16945cf36686227536decd2302c372be86da37faca68'), ('cuda_cccl/linux-x86_64/cuda_cccl-linux-x86_64-12.8.90-archive.tar.xz', '0740e9e01e4f15e17c5ab8d68bba4f8ec0eb6b84edccba4ac45112d2d2174e4b'), ) +EIGEN_REVISION = '21e4582d1739107337a03460c81412981130373e' +EIGEN_SHA256 = '399bbf2c2239b286f21edfebfe1430a689c0995acd32e8017701346f6d35be7a' +EIGEN_ARCHIVE = ROOT / 'deployment' / 'vendor' / f'eigen-{EIGEN_REVISION}.tar.gz' +EIGEN_TRELLIS_REVISION = '5565d240c4a494caaf9ece7a554542b76ffa36d3' + + +def startup_resources(tmp=Path('/tmp'), cgroup=Path('/sys/fs/cgroup')): + """Small capacity diagnostics, without listing environment or credentials.""" + disk = shutil.disk_usage(tmp) + limits = [] + for path in (cgroup / 'memory.max', cgroup / 'memory' / 'memory.limit_in_bytes'): + try: + value = path.read_text().strip() + if value.isdecimal() and 0 < int(value) < 2**60: + limits.append(int(value)) + except OSError: + pass + return {'tmp_disk_total_bytes': disk.total, 'tmp_disk_free_bytes': disk.free, + 'cgroup_memory_limit_bytes': min(limits) if limits else None} + + +def prepare_eigen(checkout, archive=None): + """Unpack the exact TRELLIS submodule from a checksum-pinned local archive.""" + checkout = Path(checkout).absolute() + archive = Path(archive) if archive is not None else EIGEN_ARCHIVE + if hashlib.sha256(archive.read_bytes()).hexdigest() != EIGEN_SHA256: + raise RuntimeError('Vendored Eigen archive checksum mismatch') + target = checkout / 'o-voxel' / 'third_party' / 'eigen' + if any(path.is_symlink() for path in (checkout, target, *target.parents) + if path == checkout or checkout in path.parents): + raise RuntimeError('Eigen extraction path must not contain symlinks') + boundary = target.resolve() + if not boundary.is_relative_to(checkout.resolve()): + raise RuntimeError('Eigen extraction escaped its checkout') + with tarfile.open(archive, 'r:gz') as source: + entries = [] + for member in source.getmembers(): + relative = PurePosixPath(member.name) + if (relative.is_absolute() or not relative.parts or + relative.parts[0] != f'eigen-{EIGEN_REVISION}' or + '..' in relative.parts or '\\' in member.name or ':' in member.name or + not (member.isdir() or member.isfile())): + raise RuntimeError('Unsafe member in vendored Eigen archive') + destination = target.joinpath(*relative.parts[1:]) + if (not destination.resolve().is_relative_to(boundary) or + any(path.is_symlink() for path in (destination, *destination.parents) + if path == target or target in path.parents)): + raise RuntimeError('Unsafe destination for vendored Eigen archive') + if destination.exists() and destination.is_file() and destination.stat().st_nlink > 1: + raise RuntimeError('Eigen extraction must not overwrite a hard-linked file') + entries.append((member, destination)) + # Validate every member before writing any headers. Tar links, special + # files and traversal paths are never materialized. + for member, destination in entries: + if member.isdir(): + destination.mkdir(parents=True, exist_ok=True) + else: + destination.parent.mkdir(parents=True, exist_ok=True) + with source.extractfile(member) as reader, destination.open('wb') as writer: + shutil.copyfileobj(reader, writer) + if not (target / 'Eigen' / 'Core').is_file() or not (target / 'COPYING.MPL2').is_file(): + raise RuntimeError('Vendored Eigen headers or license are missing') + return target + + +def prepare_submodules(name, revision, checkout, env): + if name == 'TRELLIS.2': + if revision != EIGEN_TRELLIS_REVISION: + raise RuntimeError('Review the vendored Eigen revision before updating TRELLIS.2') + prepare_eigen(checkout) + else: + run(['git', '-C', checkout, 'submodule', 'update', '--init', '--recursive', '--depth=1'], env) def status(): @@ -82,13 +177,23 @@ def run(args, env=None, cwd=None): process.wait(timeout=5) raise if code: - raise RuntimeError(f'Command failed (exit {process.returncode}): {args[0]} {args[1:4]}') + # Read only a bounded tail even when verbose compiler output is + # enormous. Public status responses never include this error text. + tail = '' + try: + with LOG.open('rb') as output: + output.seek(0, os.SEEK_END) + output.seek(max(0, output.tell() - 16000)) + tail = output.read(16000).decode('utf-8', errors='replace')[-4000:] + except OSError: + pass + raise RuntimeError(f'Command failed (exit {process.returncode}): {args[0]} {args[1:4]}\n{tail}') def runtime_environment(): """Publish only known settings, without waiting on storage or importing Torch.""" env = os.environ.copy() - env.update(CUDA_HOME=str(WORK / 'cuda-12.8'), BUILD_TARGET='cuda', TORCH_CUDA_ARCH_LIST='8.9+PTX', + env.update(CUDA_HOME=str(WORK / 'cuda-12.8'), BUILD_TARGET='cuda', TORCH_CUDA_ARCH_LIST=cuda_architectures(), MAX_JOBS='2', CMAKE_BUILD_PARALLEL_LEVEL='2', ATTN_BACKEND='sdpa', SPARSE_ATTN_BACKEND='sdpa', TORCH_COMPILE_DISABLE='1', SUPRATEX_RUNTIME_PROFILE='torch28', FLEX_GEMM_USE_AUTOTUNE_CACHE='0', FLEX_GEMM_AUTOSAVE_AUTOTUNE_CACHE='0', @@ -179,11 +284,12 @@ def copy_atomic(source, target): def build(): try: WORK.mkdir(parents=True, exist_ok=True) + update(stage='preparing_cuda', resources=startup_resources()) import torch assert torch.__version__ == '2.8.0+cu128', torch.__version__ assert torch.version.cuda == '12.8', torch.version.cuda update(stage='preparing_cuda', torch=torch.__version__, cuda=torch.version.cuda, - python=sys.version.split()[0], architecture='8.9+PTX', completed=[]) + python=sys.version.split()[0], architecture=cuda_architectures(), completed=[]) env = configure_cuda() run([Path(env['CUDA_HOME']) / 'bin/nvcc', '--version'], env) persistent = CACHE @@ -223,7 +329,7 @@ def build(): run(['git', '-C', checkout, 'remote', 'add', 'origin', url], env) run(['git', '-C', checkout, 'fetch', '--depth=1', 'origin', revision], env) run(['git', '-C', checkout, 'checkout', '--detach', revision], env) - run(['git', '-C', checkout, 'submodule', 'update', '--init', '--recursive', '--depth=1'], env) + prepare_submodules(name, revision, checkout, env) build_env = env.copy() build_env['NVCC_APPEND_FLAGS'] = O_VOXEL_FLAGS if subdir == 'o-voxel' else '--extended-lambda' if name == 'CuMesh' else '' run([*pip, 'wheel', '-v', '--no-build-isolation', '--no-deps', '--wheel-dir', wheels, checkout / subdir], build_env) diff --git a/hf_space/ovoxel_results.py b/hf_space/ovoxel_results.py index 2ec444a5f7d4b4a38b1f0950c4e374ff925016f1..7636574f20ffeccccedc25226d457f7a94668186 100644 --- a/hf_space/ovoxel_results.py +++ b/hf_space/ovoxel_results.py @@ -35,7 +35,8 @@ MAX_PREVIEW_BYTES = 64*1024*1024 _preview_lock = threading.Lock() LOG = logging.getLogger(__name__) REFERENCE_FILES = ('camera.json', 'condition_rgba.png', 'albedo.png', 'mra.png') -REFERENCE_OPTIONAL_FILES = ('condition_rgb.png',) +OBSERVATION_FILES = ('observation_keep.npy', 'observation_admission.json') +REFERENCE_OPTIONAL_FILES = ('condition_rgb.png', *OBSERVATION_FILES) REFERENCE_DIRECTORY = 'reference-input' LEGACY_REPROJECTION_WARNING = ( 'Reference-view reprojection is unavailable for this older result because its original ' @@ -72,6 +73,7 @@ def persist(source, report): reference = source.parent/'asset_work'/'input' reference_saved = False if reference.is_dir(): + _stage_observation_filter(source, reference) _persist_reference_inputs(reference, destination) reference_saved = True views = destination/'views' @@ -99,6 +101,39 @@ def _file_digest(path): return digest.hexdigest() +def _observation_keep(folder): + """Reuse the exact admission mask from completion, including after a restart.""" + import numpy as np + paths = [Path(folder)/name for name in OBSERVATION_FILES] + present = [path.is_file() for path in paths] + if not any(present): + return None + if not all(present) or any(path.is_symlink() for path in paths): + raise ValueError('The saved observation filter is incomplete') + if paths[0].stat().st_size > 2*1024*1024 or paths[1].stat().st_size > 64*1024: + raise ValueError('The saved observation filter is invalid') + keep = np.load(paths[0], allow_pickle=False) + if not isinstance(keep, np.ndarray) or keep.dtype != np.bool_ or keep.shape != (1024, 1024): + raise ValueError('The saved observation filter mask is invalid') + admission = json.loads(paths[1].read_text(encoding='utf-8')) + if not isinstance(admission, dict): + raise ValueError('The saved observation filter metadata is invalid') + return keep + + +def _stage_observation_filter(source, reference): + """Bind completion sidecars to the observations used by deferred GLB export.""" + sidecars = [source.with_suffix('.'+name) for name in OBSERVATION_FILES] + present = [path.is_file() for path in sidecars] + if not any(present) and not (reference/'image_depth.npy').is_file(): + return # Pre-upgrade completions have no image-depth gate. + if not all(present) or any(path.is_symlink() for path in sidecars): + raise ValueError('Completion did not save its observation filter') + for source_path, name in zip(sidecars, OBSERVATION_FILES): + shutil.copyfile(source_path, reference/name) + _observation_keep(reference) + + def _validate_reference_inputs(folder, bindings=None): """Validate the exact observed maps and camera; never substitute raw RGB.""" from PIL import Image @@ -130,6 +165,7 @@ def _validate_reference_inputs(folder, bindings=None): with Image.open(folder/'condition_rgb.png') as rgb, Image.open(folder/'condition_rgba.png') as rgba: if rgb.size != rgba.size: raise ValueError('The saved reference RGB and RGBA dimensions differ') + _observation_keep(folder) actual = {name: _file_digest(folder/name) for name in names} if bindings is not None and actual != bindings: raise ValueError('The saved reference inputs changed or failed integrity validation') @@ -409,11 +445,23 @@ def _extract(identifier, token, *, stage, stages, prepare_download=None): local_output = work/'asset.glb' # Export-time padding precedes native file writing; continuous observed # gutters follow reprojection. Keep these two real boundaries separate. - with observe_stages(lambda name: stage('export_gutter' if name == 'gutter' else name)): - _export_local(source, local_output, gutter=reference is None) + from hf_space.gpu_mode import enabled as zero_gpu_enabled + if zero_gpu_enabled(): + from hf_space.zerogpu_runtime import export + offset = time.monotonic() - started + result, gpu_stages = export(str(source), str(local_output), + str(reference) if reference is not None else None, str(work/'reproject'), token) + local_output = Path(result) + stages.extend(dict(event, start_seconds=round(offset+event['start_seconds'], 3)) + for event in gpu_stages) + else: + with observe_stages(lambda name: stage('export_gutter' if name == 'gutter' else name)): + _export_local(source, local_output, gutter=reference is None) + if reference is not None: + from supratex.inference.reprojection import run_front_view + local_output = run_front_view(local_output, reference, work/'reproject', stage=stage, + keep=_observation_keep(reference)) if reference is not None: - from supratex.inference.reprojection import run_front_view - local_output = run_front_view(local_output, reference, work/'reproject', stage=stage) postprocess = {'reprojection':'applied', 'gutter':'observation_guided', 'postprocess_warning':None} else: diff --git a/hf_space/preview_recovery.py b/hf_space/preview_recovery.py index 77fb464f60a98c3ad913c27e6264ab0c7abe1707..8b6db611db7a549864bab07d27f78e9223d43808 100644 --- a/hf_space/preview_recovery.py +++ b/hf_space/preview_recovery.py @@ -8,6 +8,17 @@ import time def recover(identifier): + from hf_space.gpu_mode import enabled + if enabled(): + from hf_space.zerogpu_runtime import recover as gpu_recover + from hf_space import ovoxel_results + if not (ovoxel_results.directory(identifier)/'ready.json').is_file(): + gpu_recover(identifier) + return ovoxel_results.preview_html(identifier), identifier + return _recover(identifier) + + +def _recover(identifier, *, render_html=True): import numpy as np import torch from hf_space import ovoxel_results @@ -18,7 +29,7 @@ def recover(identifier): target = ovoxel_results.directory(identifier) source = Path('/data/generated/ovoxel-recovery-v1')/identifier if (target/'ready.json').is_file(): - return ovoxel_results.preview_html(identifier), identifier + return (ovoxel_results.preview_html(identifier) if render_html else None), identifier report = json.loads((source/'report.json').read_text()) backup = json.loads((source/'backup.json').read_text()) if report.get('weight_origin')!='released' or report.get('test_only') is not False: @@ -66,4 +77,4 @@ def recover(identifier): shutil.copyfile(source/'timings.jsonl',target/'original-failed-timings.jsonl') (target/'ready.json').write_text(json.dumps({'complete':True,'identifier':identifier, 'project_revision':report['project_revision'],'output_kind':'final_ovoxel','recovered_preview':True})) - return ovoxel_results.preview_html(identifier),identifier + return (ovoxel_results.preview_html(identifier) if render_html else None),identifier diff --git a/hf_space/release_backend.py b/hf_space/release_backend.py index 126fd42e6303496350f152b41b6d466b0fc6ec7c..e7f678993eb661152b356b5ba56abbf36f07fe55 100644 --- a/hf_space/release_backend.py +++ b/hf_space/release_backend.py @@ -1,5 +1,6 @@ """Prepare the pinned release and one reusable model worker on every startup.""" from pathlib import Path +from contextlib import nullcontext import atexit import copy import json @@ -12,16 +13,20 @@ import traceback import uuid from hf_space import native_build, generation_progress +from hf_space.gpu_mode import enabled as zero_gpu_enabled ROOT=Path(__file__).resolve().parents[1] JOBS_ROOT=Path('/tmp/supratex-jobs') MODEL_REVISION='1bc02eb969a8ed0f1418c2bbe1e409a8481f834c' CODE_REVISION='e4dfbbde6a8d76e25deef5aba5d836f03994967d' +TEXTURE_FILTER_REVISION='cee4435652d0f23737bd28d5704e0fae9830b411' CHECKPOINT='denoiser_step0030000.safetensors' MODEL_ROOT=Path('/data/models') STATE=Path('/tmp/supratex-release-preparation.json') _lock=threading.Lock() _worker=None +_foreground_worker=None +_foreground_started=False _epoch=os.environ.get('SUPRATEX_STARTUP_EPOCH') or uuid.uuid4().hex # Generated binding identities include resolved cache paths. Never reuse a # previous process's prepare marker after a cache/transport upgrade. @@ -33,6 +38,7 @@ _generation_lock=threading.Lock() _state_lock=threading.Lock() _persist_lock=threading.Lock() _snapshot={'stage':'pending','epoch':_epoch} +_foreground_snapshot={'stage':'pending','epoch':_epoch} _staging_process=None _runtime_monitor=None _shutdown=threading.Event() @@ -52,6 +58,80 @@ def model_status(): with _state_lock: return copy.deepcopy(_snapshot) + +def foreground_status(): + if not zero_gpu_enabled(): + return model_status() + with _state_lock: + return copy.deepcopy(_foreground_snapshot) + + +def _update_foreground(**changes): + global _foreground_snapshot + with _state_lock: + _foreground_snapshot = _foreground_snapshot | changes | {'updated_unix': time.time()} + print('[foreground-prepare] ' + json.dumps(changes), flush=True) + + +def _stage_foreground_models(): + """Verify and mirror only the pinned RMBG group, sharing existing blobs.""" + from hf_space.download_models import source_inventory + from hf_space.weight_cache import register_sources, _mirror + spec = json.loads((ROOT/'configs/pipeline.json').read_text())['external_models']['rmbg2'] + operation = dict(spec, kind='hub', name='rmbg2', destination='external/rmbg2') + sources = source_inventory(operation, MODEL_ROOT) + total = sum(record['size'] for record in sources) + _update_foreground(startup_progress={'label':'Preparing background-removal model', + 'completed':0, 'total':total, 'unit':'bytes'}) + copied = {} + def progress(source, completed, size): + copied[str(source)] = completed + _update_foreground(startup_progress={'label':'Preparing background-removal model', + 'completed':sum(copied.values()), 'total':total, 'unit':'bytes'}) + def retry(source, attempt, attempts): + _update_foreground(startup_progress={'label':f'Reconnecting to model storage (attempt {attempt} of {attempts})', + 'completed':sum(copied.values()), 'total':total, 'unit':'bytes'}) + register_sources(sources, retry=retry) + destination = Path('/tmp/supratex-foreground-models-v1') + _mirror(MODEL_ROOT/'external/rmbg2', destination, progress) + return destination + + +def prepare_foreground_models(): + """CPU-only, independent of native compilation and the 18-model pipeline.""" + try: + _update_foreground(stage='staging_local_weights', error=None, startup_progress={ + 'label':'Checking background-removal model files', 'completed':0, 'total':1, 'unit':'steps'}) + model_path = _stage_foreground_models() + _update_foreground(stage='loading_models', startup_progress={ + 'label':'Loading background-removal model', 'completed':0, 'total':1, 'unit':'models'}) + from hf_space import zerogpu_runtime + def notify(label, completed, total): + _update_foreground(startup_progress={'label':label, 'completed':completed, + 'total':total, 'unit':'models'}) + summary = zerogpu_runtime.preload_foreground(model_path, notify) + _update_foreground(stage='ready_released', error=None, resident_models=summary, startup_progress={ + 'label':'Background-removal model ready', 'completed':1, 'total':1, 'unit':'models'}) + except BaseException: + _update_foreground(stage='failed', error=traceback.format_exc()[-8000:], startup_progress={ + 'label':'Background-removal model preparation failed. Please try again later.', + 'completed':0, 'total':1, 'unit':'steps'}) + traceback.print_exc() + + +def start_foreground_preparation(): + global _foreground_worker, _foreground_started + if not zero_gpu_enabled(): + return start_preparation() + with _lock: + if not _foreground_started: + _foreground_started = True + _foreground_worker = threading.Thread(target=prepare_foreground_models, daemon=True, + name='foreground-model-prepare') + _foreground_worker.start() + return _foreground_worker + + def update(**changes): global _snapshot with _persist_lock: @@ -87,10 +167,10 @@ def _stage_models(): # longer rehashes or copies large checkpoints through Bucket FUSE. sources.extend(download(operation,MODEL_ROOT,inventory_only=True)) from hf_space.weight_cache import mirror_models,mirror_project,register_sources + total=sum(record['size'] for record in sources) update(stage='staging_local_weights',startup_progress={ - 'label':'Preparing local weight files','completed':0,'total':0}) + 'label':'Preparing local weight files','completed':0,'total':total,'unit':'bytes'}) project_root=MODEL_ROOT/'project'/revision - total=sum(record['size'] for record in sources) copied={} def file_progress(source,completed,size): copied[str(source)]=completed @@ -115,7 +195,7 @@ def _stage_models(): project_revision=revision,code_revision=CODE_REVISION,weight_origin='released', actual_training_steps=30000,random_weights_present=False, checkpoint=CHECKPOINT,model_loading='startup_pool_with_stage_offload', - startup_progress={'label':'Initializing model runtime','completed':0,'total':0}) + startup_progress={'label':'Initializing model runtime','completed':0,'total':1,'unit':'steps'}) return config except BaseException: update(stage='failed',error=traceback.format_exc()) @@ -191,7 +271,8 @@ def prepare_models(): from hf_space.startup_guard import StartupCleanupError _startup_deadline=time.monotonic()+STARTUP_TOTAL_SECONDS try: - update(stage='waiting_native_build',startup_progress={'label':'Preparing runtime','completed':0,'total':0}) + update(stage='waiting_native_build',startup_progress={ + 'label':'Preparing runtime','completed':0,'total':1,'unit':'steps'}) native_build.start() deadline=time.monotonic()+_remaining(NATIVE_WAIT_SECONDS) while native_build.status().get('stage') not in ('compiled','failed'): @@ -204,10 +285,11 @@ def prepare_models(): try: _remaining(STARTUP_TOTAL_SECONDS) update(stage='staging_local_weights',attempt=attempt,attempts=STARTUP_ATTEMPTS, - error=None,startup_progress={'label':'Preparing local weight files','completed':0,'total':0}) + error=None,startup_progress={ + 'label':'Checking model files','completed':0,'total':1,'unit':'steps'}) config=_stage_attempt(attempt) update(stage='loading_models',startup_progress={ - 'label':'Initializing model runtime','completed':0,'total':0}) + 'label':'Initializing model runtime','completed':0,'total':1,'unit':'steps'}) launch_runtime(config) update(stage='ready_released',error=None) _start_runtime_monitor() @@ -220,11 +302,11 @@ def prepare_models(): if attempt==STARTUP_ATTEMPTS: raise update(stage='retrying',error=error,startup_progress={ - 'label':'Retrying model startup','completed':0,'total':0}) + 'label':'Retrying model startup','completed':0,'total':1,'unit':'steps'}) time.sleep(min(STARTUP_BACKOFF_SECONDS*attempt,_remaining(STARTUP_TOTAL_SECONDS))) except BaseException: update(stage='failed',error=traceback.format_exc()[-8000:],startup_progress={ - 'label':'Model startup failed. Please try again later.','completed':0,'total':0}) + 'label':'Model startup failed. Please try again later.','completed':0,'total':1,'unit':'steps'}) traceback.print_exc() def start_preparation(): @@ -239,6 +321,10 @@ def start_preparation(): def stop_runtime(): global _runtime_process + if zero_gpu_enabled(): + runtime = sys.modules.get('hf_space.zerogpu_runtime') + if runtime is not None: + runtime.clear() process=_runtime_process if process is not None: from hf_space.startup_guard import stop_process @@ -260,6 +346,17 @@ atexit.register(stop_preparation) def launch_runtime(config): global _runtime_process + if zero_gpu_enabled(): + from hf_space import zerogpu_runtime + def notify(label, completed, total): + update(startup_progress={'label': label, 'completed': completed, + 'total': total, 'unit': 'models'}) + summary = zerogpu_runtime.preload(config, notify) + update(resident_models=summary, model_loading='cpu_pool_with_per_user_zerogpu', + startup_progress={'label': 'Models ready', + 'completed': summary['loaded_models'], + 'total': summary['loaded_models'], 'unit': 'models'}) + return from hf_space.startup_guard import process_options if _runtime_process is not None: stop_runtime() @@ -302,6 +399,9 @@ def launch_runtime(config): def _start_runtime_monitor(): global _runtime_monitor + if zero_gpu_enabled(): + # There is deliberately no resident CUDA subprocess to monitor. + return process=_runtime_process def monitor(): last=None @@ -335,14 +435,27 @@ def ensure_prepared(generation_token=None): def run_resident(job, seed, timeout_seconds, generation_token, *, operation='generate', foreground_prepared=False): - with _generation_lock: - if model_status().get('stage')!='ready_released': + zero_gpu = zero_gpu_enabled() + # Each ZeroGPU lease owns an isolated fork. Only the dedicated runtime's + # single stdin/result protocol needs a process-wide execution lock. + with nullcontext() if zero_gpu else _generation_lock: + if operation not in ('generate','remove_background','generate_geometry','generate_texture') or type(foreground_prepared) is not bool: + raise ValueError('Invalid internal model operation') + state = foreground_status() if zero_gpu and operation == 'remove_background' else model_status() + if state.get('stage')!='ready_released': raise RuntimeError('Background model loading is incomplete') + if zero_gpu: + from hf_space import zerogpu_runtime + if operation == 'remove_background': + return zerogpu_runtime.foreground(str(job)) + if operation == 'generate_geometry': + return zerogpu_runtime.geometry(str(job), int(seed), foreground_prepared) + if operation == 'generate_texture': + return zerogpu_runtime.texture(str(job), int(seed)) + return zerogpu_runtime.generate(str(job), int(seed), foreground_prepared) process=_runtime_process if process is None or process.poll() is not None: raise RuntimeError('The preloaded model worker is unavailable') - if operation not in ('generate','remove_background') or type(foreground_prepared) is not bool: - raise ValueError('Invalid internal model operation') request={'epoch':_epoch,'job':str(job),'seed':int(seed), 'operation':operation,'foreground_prepared':foreground_prepared} process.stdin.write(json.dumps(request)+'\n') @@ -365,7 +478,7 @@ def run_resident(job, seed, timeout_seconds, generation_token, *, operation='gen generation_progress.update(generation_token,label={ 'frontend':'Generating geometry and reference materials', 'completion':'Completing surface materials and rendering o-voxel previews', - }.get(stage,'Running the latest single-view pipeline')) + }.get(stage,'Starting generation…')) time.sleep(.5) response=json.loads(response_path.read_text()) if response.get('epoch')!=_epoch: @@ -380,7 +493,7 @@ def run_resident(job, seed, timeout_seconds, generation_token, *, operation='gen def remove_background(image,timeout_seconds=120): """Run the queued cutout in the one startup-loaded, serial model worker.""" - if model_status().get('stage')!='ready_released': + if foreground_status().get('stage')!='ready_released': raise RuntimeError('Model files are not ready') job=JOBS_ROOT/uuid.uuid4().hex job.mkdir(parents=True) @@ -415,8 +528,8 @@ def generate(image,seed,progress=None,timeout_seconds=1800,generation_token=None output=job/'asset.npz' started=time.monotonic() generation_progress.update(generation_token,status='running',job=identifier, - generation_started_unix=time.time(),label='Running the latest single-view pipeline') - if progress:progress(.05,'Running the latest single-view pipeline') + generation_started_unix=time.time(),label='Starting generation…') + if progress:progress(.05,'Starting generation…') resident_models=run_resident(job,seed,timeout_seconds,generation_token,foreground_prepared=foreground_prepared) receipt_path=job/'asset_work/run.json' receipt=json.loads(receipt_path.read_text()) if receipt_path.exists() else {} @@ -438,7 +551,8 @@ def generate(image,seed,progress=None,timeout_seconds=1800,generation_token=None report={'status':'complete','asset_id':identifier,'code_revision':CODE_REVISION, 'project_revision':MODEL_REVISION,'checkpoint':CHECKPOINT,'weight_origin':'released', 'elapsed_seconds':time.monotonic()-started,'output_kind':'final_ovoxel','backview':False, - 'refine':False,'model_loading':'startup_pool_with_stage_offload','resident_models':resident_models,'run':receipt, + 'refine':False,'model_loading':('cpu_pool_with_per_user_zerogpu' if zero_gpu_enabled() + else 'startup_pool_with_stage_offload'),'resident_models':resident_models,'run':receipt, 'preview':preview,'voxel_count':voxel_count,'automatic_glb_export':False, 'foreground_prepared':foreground_prepared} (job/'report.json').write_text(json.dumps(report,indent=2)) @@ -448,6 +562,90 @@ def generate(image,seed,progress=None,timeout_seconds=1800,generation_token=None return str(output),str(output),report +def _split_job(generation_token, label): + if model_status().get('stage') != 'ready_released': + raise RuntimeError('Model files are not ready') + job = JOBS_ROOT / uuid.uuid4().hex + job.mkdir(parents=True) + generation_progress.update(generation_token, status='running', job=job.name, + generation_started_unix=time.time(), label=label) + return job + + +def _split_result(job, started, resident_models, *, geometry=False): + import numpy as np + result = json.loads((job / 'stage-result.json').read_text(encoding='utf-8')) + expected = 'geometry' if geometry else 'final_ovoxel' + if result.get('status') != 'complete' or result.get('output_kind') != expected: + raise RuntimeError('The generation stage did not complete') + output = job / ('geometry.npz' if geometry else 'asset.npz') + with np.load(output, allow_pickle=False) as asset: + required = {'vertices', 'faces'} if geometry else {'vertices','faces','coords','attrs','resolution','origin','voxel_size'} + if not required <= set(asset.files) or not len(asset['vertices']) or not len(asset['faces']): + raise RuntimeError('The generated geometry is incomplete') + if not geometry and (not len(asset['coords']) or asset['attrs'].shape[1] != 6 + or not np.isfinite(asset['attrs']).all()): + raise RuntimeError('The generated material is incomplete') + face_count = len(asset['faces']) + voxel_count = None if geometry else len(asset['coords']) + if geometry: + from hf_space.split_pipeline import CHECKPOINT_DIRECTORY, validate_checkpoint + validate_checkpoint(job / CHECKPOINT_DIRECTORY) + preview = json.loads((job / 'preview.json').read_text(encoding='utf-8')) + expected_frames = {f'{mode}-{index:02d}.webp' for mode in preview['modes'] + for index in range(preview['views'])} + frames = list((job / 'views').glob('*.webp')) + if not expected_frames or {path.name for path in frames} != expected_frames or any(not p.stat().st_size for p in frames): + raise RuntimeError('The preview render set is incomplete') + if geometry and (preview['modes'] != ['clay'] or preview.get('default_mode') != 'clay'): + raise RuntimeError('The geometry preview must contain clay views only') + if list(job.rglob('*.glb')) or preview.get('automatic_glb_export') is not False: + raise RuntimeError('Generation must not export a mesh before selection') + report = {'status':'complete', 'asset_id':job.name, 'code_revision':CODE_REVISION, + 'project_revision':MODEL_REVISION, 'checkpoint':CHECKPOINT, 'weight_origin':'released', + 'elapsed_seconds':time.monotonic()-started, 'output_kind':expected, + 'backview':False, 'refine':False, 'automatic_glb_export':False, + 'model_loading':('cpu_pool_with_per_user_zerogpu' if zero_gpu_enabled() + else 'startup_pool_with_stage_offload'), + 'resident_models':resident_models, 'run':result, 'preview':preview, + 'face_count':face_count, 'geometry_reused':not geometry} + if voxel_count is not None: + report['voxel_count'] = voxel_count + admission = output.with_suffix('.observation_admission.json') + report['depth_pre_filter'] = { + 'code_revision': TEXTURE_FILTER_REVISION, + 'admission': json.loads(admission.read_text(encoding='utf-8')) if admission.is_file() else None, + } + (job / 'report.json').write_text(json.dumps(report, indent=2), encoding='utf-8') + if not (job / 'timings.jsonl').exists(): + (job / 'timings.jsonl').write_text('', encoding='utf-8') + return str(output), str(output), report + + +def generate_geometry(image, seed, generation_token=None, *, foreground_prepared=False, timeout_seconds=1800): + if type(foreground_prepared) is not bool: + raise ValueError('Prepared foreground must be an internal Boolean') + started = time.monotonic() + job = _split_job(generation_token, 'Generating geometry') + image.save(job / 'input.png') + summary = run_resident(job, int(seed), timeout_seconds, generation_token, + operation='generate_geometry', foreground_prepared=foreground_prepared) + return _split_result(job, started, summary, geometry=True) + + +def generate_texture(geometry_folder, seed, generation_token=None, *, timeout_seconds=1800): + import shutil + from hf_space.split_pipeline import CHECKPOINT_DIRECTORY, copy_checkpoint + geometry_folder = Path(geometry_folder) + started = time.monotonic() + job = _split_job(generation_token, 'Generating materials') + # Copy from the verified durable checkpoint before acquiring a GPU lease. + copy_checkpoint(geometry_folder / CHECKPOINT_DIRECTORY, job / CHECKPOINT_DIRECTORY, seed=seed) + shutil.copyfile(geometry_folder / 'input.png', job / 'input.png') + summary = run_resident(job, int(seed), timeout_seconds, generation_token, operation='generate_texture') + return _split_result(job, started, summary) + + if __name__=='__main__': if len(sys.argv)!=3 or sys.argv[1]!='--stage': raise SystemExit('This module is an internal model preparation worker') diff --git a/hf_space/resident_pipeline.py b/hf_space/resident_pipeline.py index f523c1df2ac85760cd47f274be2011c56a559783..60180946609b049920b549572659a2fd2ee9a4ee 100644 --- a/hf_space/resident_pipeline.py +++ b/hf_space/resident_pipeline.py @@ -12,6 +12,32 @@ from hf_space.model_pool import ModelPool ROOT = Path(__file__).resolve().parents[1] +def _configure_preload_threads(torch, *, cpu_only): + # spaces 0.51.3 already sets the process-wide inter-op count at import. + # PyTorch permits that setter only once, even before the pool is consumed. + # Intra-op is adjustable and stays serial for the later ZeroGPU fork. + torch.set_num_threads(1 if cpu_only else 4) + if not cpu_only and torch.get_num_interop_threads() != 1: + torch.set_num_interop_threads(1) + + +class ForegroundResident: + """The standalone cutout task needs only RMBG, without native 3D modules.""" + def __init__(self, model_path): + self.config = {'rmbg_model': str(model_path)} + self.pool = ModelPool(directory='/tmp/supratex-foreground-state-v1') + + def prepare(self, notify): + import torch + from transformers import AutoModelForImageSegmentation + _configure_preload_threads(torch, cpu_only=True) + notify('Loading background-removal model', 0, 1) + self.pool.add('background', AutoModelForImageSegmentation.from_pretrained( + self.config['rmbg_model'], trust_remote_code=True, local_files_only=True).eval()) + self.pool.sealed = True + notify('Background-removal model ready', 1, 1) + + class ResidentPipeline: def __init__(self, config): self.config_path = str(config) @@ -23,29 +49,31 @@ class ResidentPipeline: self.material_step = None self.material_processor = None - def prepare(self, notify): + def prepare(self, notify, *, cpu_only=False): import torch from scripts.run_example import read_config from supratex.inference.runtime import configure_inference_runtime # Startup constructors run on host tensors; no model forward is called. config = read_config(Path(self.config)) self.config = config - os.environ.update(TORCH_HOME=config['torch_home'], HF_HUB_OFFLINE='1', - TRANSFORMERS_OFFLINE='1', TOKENIZERS_PARALLELISM='false', + os.environ.update(TORCH_HOME=config['torch_home'], TOKENIZERS_PARALLELISM='false', ATTN_BACKEND='sdpa', SPARSE_ATTN_BACKEND='sdpa', SUPRATEX_RUNTIME_PROFILE='torch28', SUPRATEX_LOW_VRAM='1') - # The Space has eight CPU cores. Bounded constructor threads leave - # headroom for uploads and avoid oversubscribing the container quota. - torch.set_num_threads(4) - torch.set_num_interop_threads(1) - if not torch.cuda.is_available(): + # The isolated dedicated-GPU worker can globally disable Hub access. + # ZeroGPU preloads in the HTTP process, which still serves OAuth and + # Bucket requests; local model paths and local_files_only are sufficient. + if not cpu_only: + os.environ.update(HF_HUB_OFFLINE='1', TRANSFORMERS_OFFLINE='1') + _configure_preload_threads(torch, cpu_only=cpu_only) + if not cpu_only and not torch.cuda.is_available(): raise RuntimeError('The GPU is unavailable; startup models cannot become ready') for path in (ROOT, ROOT/'deployment/runtime/moge', ROOT/'deployment/runtime/utils3d', Path(config['geometry_package']), Path(config['qwen_runtime']), Path(config['qwen_deps']) if config.get('qwen_deps') else None): if path is not None and str(path) not in sys.path: sys.path.insert(0, str(path)) - configure_inference_runtime(attention='sdpa', low_vram=True, device='cuda', + configure_inference_runtime(attention='sdpa', low_vram=True, + device='cpu' if cpu_only else 'cuda', material_backend='qwen-image') total = 18 done = 0 @@ -129,10 +157,11 @@ class ResidentPipeline: build_completion_model(c, load_completion_state_dict(checkpoint, torch_module=torch), assign_weights=True).eval()) # Establish a working CUDA context before enabling Generate. Constructors # above load host weights; they do not need neural warm-up forwards. - probe = torch.zeros(1, device='cuda') - torch.cuda.synchronize() - del probe - torch.cuda.empty_cache() + if not cpu_only: + probe = torch.zeros(1, device='cuda') + torch.cuda.synchronize() + del probe + torch.cuda.empty_cache() self.pool.sealed = True notify('Models ready', total, total) print('[resident-models] '+json.dumps(self.pool.summary()), flush=True) diff --git a/hf_space/runtime_worker.py b/hf_space/runtime_worker.py index 93ec17c0b15bd2a2181772a8bd48a875137c216d..43e25e1c70d1f09448287d13f4f610feff99e715 100644 --- a/hf_space/runtime_worker.py +++ b/hf_space/runtime_worker.py @@ -25,7 +25,8 @@ def atomic_json(path, value): temporary.replace(path) -def serve(config_path, status_path, epoch, stream, *, factory=None, run_job=None, run_foreground=None): +def serve(config_path, status_path, epoch, stream, *, factory=None, run_job=None, run_foreground=None, + run_geometry=None, run_texture=None): """Injectable protocol boundary; production models stay in this process.""" def status(stage, **fields): atomic_json(status_path, {"epoch": epoch, "stage": stage, @@ -68,7 +69,7 @@ def serve(config_path, status_path, epoch, stream, *, factory=None, run_job=None if request.get("epoch") != epoch: raise ValueError("Job belongs to another runtime") operation = request.get('operation', 'generate') - if operation not in ('generate', 'remove_background'): + if operation not in ('generate', 'remove_background', 'generate_geometry', 'generate_texture'): raise ValueError('Unknown internal model operation') foreground_prepared = request.get('foreground_prepared', False) if type(foreground_prepared) is not bool: @@ -81,6 +82,11 @@ def serve(config_path, status_path, epoch, stream, *, factory=None, run_job=None try: if operation == 'remove_background': (run_foreground or execute_foreground)(resident, job) + elif operation == 'generate_geometry': + (run_geometry or execute_geometry)(resident, job, int(request['seed']), + foreground_prepared=foreground_prepared) + elif operation == 'generate_texture': + (run_texture or execute_texture)(resident, job, int(request['seed'])) else: options = {'foreground_prepared': True} if foreground_prepared else {} (run_job or execute)(resident, job, int(request["seed"]), **options) @@ -121,7 +127,7 @@ def execute_foreground(resident, job): atomic_json(job / 'foreground.json', metadata) -def execute(resident, job, seed, *, foreground_prepared=False): +def _arguments(resident, job, seed): from scripts.run_example import parse_args, build_command arguments, extra = parse_args([ "--config", str(resident.config_path), "--input", str(job / "input.png"), @@ -134,10 +140,29 @@ def execute(resident, job, seed, *, foreground_prepared=False): spec = importlib.util.spec_from_file_location("_supratex_resident_infer", ROOT / "infer.py") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) - module.run(module.parse_args(command[2:]), resident=resident, + return module, module.parse_args(command[2:]) + + +def execute(resident, job, seed, *, foreground_prepared=False): + module, arguments = _arguments(resident, job, seed) + module.run(arguments, resident=resident, foreground_prepared=foreground_prepared) +def execute_geometry(resident, job, seed, *, foreground_prepared=False): + from hf_space.split_pipeline import run_geometry + _module, arguments = _arguments(resident, job, seed) + result = run_geometry(arguments, resident=resident, foreground_prepared=foreground_prepared) + atomic_json(job / 'stage-result.json', dict(result, status='complete')) + + +def execute_texture(resident, job, seed): + from hf_space.split_pipeline import run_texture + _module, arguments = _arguments(resident, job, seed) + result = run_texture(arguments, resident=resident) + atomic_json(job / 'stage-result.json', dict(result, status='complete')) + + def main(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) diff --git a/hf_space/split_pipeline.py b/hf_space/split_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..027b1391f14d3ba6273fbcf6ec9fae63729133a4 --- /dev/null +++ b/hf_space/split_pipeline.py @@ -0,0 +1,146 @@ +"""Two explicit inference requests sharing an immutable geometry checkpoint. + +The checkpoint contains the exact prepared condition, optimized camera and mesh. +Texture generation consumes those files without invoking the image/geometry +frontend. The history layer owns durable copies and their content checksums. +""" +from __future__ import annotations + +import json +from pathlib import Path +import shutil + +from hf_space.timing import span + +CHECKPOINT_DIRECTORY = 'geometry-checkpoint' +CHECKPOINT_FORMAT = 'supratex.geometry-checkpoint.v1' +CHECKPOINT_FILES = ('condition_rgba.png', 'condition_rgb.png', 'camera.json', + 'geometry.pt', 'geometry_receipt.json', 'checkpoint.json') +OPTIONAL_CHECKPOINT_FILES = ('image_depth.npy',) + + +def validate_checkpoint(folder, *, seed=None): + folder = Path(folder) + if folder.is_symlink() or not folder.is_dir(): + raise ValueError('The saved geometry checkpoint is unavailable') + for name in CHECKPOINT_FILES: + path = folder / name + if path.is_symlink() or not path.is_file() or path.stat().st_size <= 0: + raise ValueError('The saved geometry checkpoint is incomplete') + metadata = json.loads((folder / 'checkpoint.json').read_text(encoding='utf-8')) + declared = metadata.get('files') + if (metadata.get('format') != CHECKPOINT_FORMAT + or type(metadata.get('seed')) is not int + or type(metadata.get('remove_floaters_applied')) is not bool + or declared not in (list(CHECKPOINT_FILES[:-1]), + [*CHECKPOINT_FILES[:-1], *OPTIONAL_CHECKPOINT_FILES])): + raise ValueError('The saved geometry checkpoint has an unsupported format') + for name in OPTIONAL_CHECKPOINT_FILES: + path = folder / name + if name in declared: + if path.is_symlink() or not path.is_file() or path.stat().st_size <= 0: + raise ValueError('The saved geometry checkpoint is incomplete') + elif path.exists() or path.is_symlink(): + raise ValueError('The saved geometry checkpoint contains undeclared depth') + if seed is not None and metadata['seed'] != int(seed): + raise ValueError('The saved geometry belongs to a different seed') + return metadata + + +def copy_checkpoint(source, destination, *, seed=None): + """Copy only the bounded, relative checkpoint contract into a fresh folder.""" + metadata = validate_checkpoint(source, seed=seed) + destination = Path(destination) + destination.mkdir(parents=True, exist_ok=False) + for name in (*metadata['files'], 'checkpoint.json'): + shutil.copyfile(Path(source) / name, destination / name) + return metadata + + +def _frontend_options(args): + return dict(moge_model=args.moge_model, rmbg_model=args.rmbg_model, + use_input_alpha=args.use_input_alpha, resolution_level=args.resolution_level, + geometry_package=args.geometry_package.expanduser().resolve(), geometry_model=args.geometry_model, + geometry_revision=args.geometry_revision, geometry_seed=args.seed, + geometry_initial_mesh_scale=args.geometry_initial_mesh_scale, + geometry_scale_search=args.geometry_scale_search, geometry_target_padding=args.geometry_target_padding, + geometry_max_iterations=args.geometry_max_iterations, + material_backend=args.material_backend, material_base=args.material_base, + qwen_runtime=args.qwen_runtime, qwen_deps=args.qwen_deps, + material_model=args.material_model, revision=args.material_revision, + cache_dir=args.cache_dir, local_files_only=args.local_files_only, + export_image_depth=not (getattr(args, 'no_filter', False) or getattr(args, 'no_image_depth', False))) + + +def run_geometry(args, *, resident, foreground_prepared=False): + """Run foreground/camera/geometry once; never predict material maps.""" + import numpy as np + import torch + from PIL import Image + from supratex.frontend.pipeline import run as frontend + from supratex.inference.runtime import configure_inference_runtime + from hf_space.geometry_preview import render_geometry + + job = Path(args.output).parent + checkpoint = job / CHECKPOINT_DIRECTORY + configure_inference_runtime(attention='sdpa', low_vram=True, material_backend='qwen-image') + frontend(Path(args.input), checkpoint, **_frontend_options(args), resident=resident, + foreground_prepared=foreground_prepared, geometry_only=True) + geometry = torch.load(checkpoint / 'geometry.pt', map_location='cpu', weights_only=True) + vertices, faces = geometry['verts'].float().contiguous(), geometry['faces'].long().contiguous() + camera = json.loads((checkpoint / 'camera.json').read_text(encoding='utf-8')) + cleanup = None + if args.remove_floaters: + from supratex.inference.geometry_cleanup import remove_floaters + with span('geometry.cleanup'), Image.open(checkpoint / 'condition_rgba.png') as image: + vertices, faces, cleanup = remove_floaters(vertices, faces, camera, + np.asarray(image.getchannel('A'))) + if (vertices.ndim != 2 or vertices.shape[1] != 3 or faces.ndim != 2 or faces.shape[1] != 3 + or not len(vertices) or not len(faces) or not torch.isfinite(vertices).all() + or faces.min() < 0 or faces.max() >= len(vertices)): + raise ValueError('Geometry generation returned an invalid mesh') + vertices, faces = vertices.cpu().contiguous(), faces.cpu().contiguous() + with span('geometry.save'): + # Apply the original completion cleanup once, before displaying the mesh. + # Stage two must use this exact topology, without repeating cleanup. + torch.save(dict(geometry, verts=vertices, faces=faces), checkpoint / 'geometry.pt') + np.savez(job / 'geometry.npz', vertices=vertices.numpy(), faces=faces.numpy()) + metadata = {'format': CHECKPOINT_FORMAT, 'seed': int(args.seed), + 'remove_floaters_applied': bool(args.remove_floaters), + 'files': list(CHECKPOINT_FILES[:-1])} + metadata['files'].extend(name for name in OPTIONAL_CHECKPOINT_FILES if (checkpoint / name).is_file()) + (checkpoint / 'checkpoint.json').write_text(json.dumps(metadata), encoding='utf-8') + (job / 'input-camera.json').write_text(json.dumps(camera), encoding='utf-8') + render_geometry(job, vertices, faces, camera) + return {'output_kind': 'geometry', 'face_count': len(faces), 'vertex_count': len(vertices), + 'geometry_checkpoint': CHECKPOINT_DIRECTORY, 'geometry_cleanup': cleanup, + 'foreground_prepared': bool(foreground_prepared)} + + +def run_texture(args, *, resident): + """Resume from a saved geometry handoff; there is no geometry code path.""" + from supratex.frontend.pipeline import predict_material_maps + from supratex.inference import run_completion + from supratex.inference.runtime import configure_inference_runtime + + job = Path(args.output).parent + handoff = job / 'asset_work' / 'input' + copy_checkpoint(job / CHECKPOINT_DIRECTORY, handoff, seed=args.seed) + configure_inference_runtime(attention='sdpa', low_vram=True, material_backend='qwen-image') + with span('frontend.material'): + albedo, mra = predict_material_maps( + handoff / 'condition_rgba.png', handoff / 'condition_rgb.png', handoff / 'material', + material_backend='qwen-image', material_base=args.material_base, + material_model=args.material_model, qwen_runtime=args.qwen_runtime, + qwen_deps=args.qwen_deps, revision=args.material_revision, + cache_dir=args.cache_dir, local_files_only=True, resident=resident) + shutil.copyfile(albedo, handoff / 'albedo.png') + shutil.copyfile(mra, handoff / 'mra.png') + with span('texture.completion'): + run_completion(input_dir=handoff, output=Path(args.output), model_config=args.model_config, + checkpoint=args.checkpoint, trellis_root=args.trellis_root, seed=int(args.seed), + texture_size=args.texture_size, decimation_target=args.decimation_target, + texture_format=args.texture_format, gutter=args.gutter, + remove_floaters=False, preview_only=True, resident=resident) + return {'output_kind': 'final_ovoxel', 'geometry_reused': True, + 'geometry_checkpoint': CHECKPOINT_DIRECTORY} diff --git a/hf_space/vendor/gradio-client-2.1.0.js b/hf_space/vendor/gradio-client-2.1.0.js new file mode 100644 index 0000000000000000000000000000000000000000..c273ff7d35128ae70b10ce04c6dcfb0099d8cfb1 --- /dev/null +++ b/hf_space/vendor/gradio-client-2.1.0.js @@ -0,0 +1,1940 @@ +const We = "host", Te = "queue/data", Ze = "queue/join", ye = "upload", Ke = "login", V = "config", Ve = "info", Ye = "runtime", Qe = "sleeptime", Xe = "heartbeat", et = "component_server", tt = "reset", st = "cancel", nt = "app_id", Ae = "This application is currently busy. Please try again. ", L = "Connection errored out. ", N = "Could not resolve app config. ", ot = "Could not get space status. ", it = "Could not get API info. ", le = "Space metadata could not be loaded. ", rt = "Invalid URL. A full URL path is required.", at = "Not authorized to access this space. ", xe = "Invalid credentials. Could not login. ", ct = "Login credentials are required to access this space.", ut = "File system access is only available in Node.js environments", Ce = "Root URL not found in client config", lt = "Error uploading file"; +async function be(e, s, t) { + try { + return (await (await fetch(`https://huggingface.co/api/spaces/${e}/jwt`, { + headers: { + Authorization: `Bearer ${s}`, + ...t ? { Cookie: t } : {} + } + })).json()).token || !1; + } catch { + return !1; + } +} +function pt(e) { + let s = {}; + return e.forEach(({ api_name: t, id: n }) => { + t && (s[t] = n); + }), s; +} +async function dt(e) { + const s = this.options.token ? { Authorization: `Bearer ${this.options.token}` } : {}; + if (s["Content-Type"] = "application/json", typeof window < "u" && window.gradio_config && location.origin !== "http://localhost:9876") { + if (window.gradio_config.current_page && (e = e.substring(0, e.lastIndexOf("/"))), window.gradio_config.dev_mode || typeof window < "u" && window?.BUILD_MODE === "dev") { + let t = ce( + e, + this.deep_link ? V + "?deep_link=" + this.deep_link : V + ); + const n = await this.fetch(t, { + headers: s, + credentials: "include" + }), o = await ve(n, !!this.options.auth); + o.root = e || o.root, window.gradio_config = { + ...o, + current_page: window.gradio_config.current_page + }; + } + return { ...window.gradio_config }; + } else if (e) { + let t = ce( + e, + this.deep_link ? V + "?deep_link=" + this.deep_link : V + ); + const n = await this.fetch(t, { + headers: s, + credentials: "include" + }), o = await ve(n, !!this.options.auth); + return o.root || (o.root = e), o; + } + throw new Error(N); +} +async function ve(e, s) { + if (e?.status === 401 && !s) { + const n = (await e.json())?.detail?.auth_message; + throw new Error(n || ct); + } else if (e?.status === 401 && s) + throw new Error(xe); + if (e?.status === 200) { + let t = await e.json(); + return t.dependencies?.forEach((n, o) => { + n.id === void 0 && (n.id = o); + }), t; + } else if (e?.status === 401) + throw new Error(at); + throw new Error(N); +} +async function ht() { + const { http_protocol: e, host: s } = await de( + this.app_reference, + this.options.token + ); + try { + if (this.options.auth) { + const t = await De( + e, + s, + this.options.auth, + this.fetch, + this.options.token + ); + t && this.set_cookies(t); + } + } catch (t) { + throw Error(t.message); + } +} +async function De(e, s, t, n, o) { + const r = new FormData(); + r.append("username", t?.[0]), r.append("password", t?.[1]); + let i = {}; + o && (i.Authorization = `Bearer ${o}`); + const a = await n(`${e}//${s}/${Ke}`, { + headers: i, + method: "POST", + body: r, + credentials: "include" + }); + if (a.status === 200) + return a.headers.get("set-cookie"); + throw a.status === 401 ? new Error(xe) : new Error(le); +} +function ae(e) { + if (e.startsWith("http")) { + const { protocol: s, host: t, pathname: n } = new URL(e); + return { + ws_protocol: s === "https:" ? "wss" : "ws", + http_protocol: s, + host: t + (n !== "/" ? n : "") + }; + } + return { + ws_protocol: "wss", + http_protocol: "https:", + host: new URL(e).host + }; +} +const Ne = (e) => { + let s = []; + return e.split(/,(?=\s*[^\s=;]+=[^\s=;]+)/).forEach((n) => { + const [o, r] = n.split(";")[0].split("="); + o && r && s.push(`${o.trim()}=${r.trim()}`); + }), s; +}, pe = /^[a-zA-Z0-9_\-\.]+\/[a-zA-Z0-9_\-\.]+$/, ft = /.*hf\.space\/{0,1}.*$/; +async function de(e, s) { + const t = {}; + s && (t.Authorization = `Bearer ${s}`); + const n = e.trim().replace(/\/$/, ""); + if (pe.test(n)) + try { + const r = (await (await fetch( + `https://huggingface.co/api/spaces/${n}/${We}`, + { headers: t } + )).json()).host; + return { + space_id: e, + ...ae(r) + }; + } catch { + throw new Error(le); + } + if (ft.test(n)) { + const { ws_protocol: o, http_protocol: r, host: i } = ae(n); + return { + space_id: i.split("/")[0].replace(".hf.space", ""), + ws_protocol: o, + http_protocol: r, + host: i + }; + } + return { + space_id: !1, + ...ae(n) + }; +} +const ce = (...e) => { + try { + return e.reduce((s, t) => (s = s.replace(/\/+$/, ""), t = t.replace(/^\/+/, ""), new URL(t, s + "/").toString())); + } catch { + throw new Error(rt); + } +}; +function _t(e, s, t) { + const n = { + named_endpoints: {}, + unnamed_endpoints: {} + }; + return Object.keys(e).forEach((o) => { + (o === "named_endpoints" || o === "unnamed_endpoints") && (n[o] = {}, Object.entries(e[o]).forEach( + ([r, { parameters: i, returns: a }]) => { + const u = s.dependencies.find( + (c) => c.api_name === r || c.api_name === r.replace("/", "") + )?.id || t[r.replace("/", "")] || -1, l = u !== -1 ? s.dependencies.find((c) => c.id == u)?.types : { generator: !1, cancel: !1 }; + if (u !== -1 && s.dependencies.find((c) => c.id == u)?.inputs?.length !== i.length) { + const c = s.dependencies.find((m) => m.id == u).inputs.map( + (m) => s.components.find((w) => w.id === m)?.type + ); + try { + c.forEach((m, w) => { + if (m === "state") { + const P = { + component: "state", + example: null, + parameter_default: null, + parameter_has_default: !0, + parameter_name: null, + hidden: !0 + }; + i.splice(w, 0, P); + } + }); + } catch (m) { + console.error(m); + } + } + const p = (c, m, w, P) => ({ + ...c, + description: gt(c?.type, w), + type: mt(c?.type, m, w, P) || "" + }); + n[o][r] = { + parameters: i.map( + (c) => p(c, c?.component, c?.serializer, "parameter") + ), + returns: a.map( + (c) => p(c, c?.component, c?.serializer, "return") + ), + type: l + }; + } + )); + }), n; +} +function mt(e, s, t, n) { + if (s === "Api") return e.type; + switch (e?.type) { + case "string": + return "string"; + case "boolean": + return "boolean"; + case "number": + return "number"; + } + if (t === "JSONSerializable" || t === "StringSerializable") + return "any"; + if (t === "ListStringSerializable") + return "string[]"; + if (s === "Image") + return n === "parameter" ? "Blob | File | Buffer" : "string"; + if (t === "FileSerializable") + return e?.type === "array" ? n === "parameter" ? "(Blob | File | Buffer)[]" : "{ name: string; data: string; size?: number; is_file?: boolean; orig_name?: string}[]" : n === "parameter" ? "Blob | File | Buffer" : "{ name: string; data: string; size?: number; is_file?: boolean; orig_name?: string}"; + if (t === "GallerySerializable") + return n === "parameter" ? "[(Blob | File | Buffer), (string | null)][]" : "[{ name: string; data: string; size?: number; is_file?: boolean; orig_name?: string}, (string | null))][]"; +} +function gt(e, s) { + return s === "GallerySerializable" ? "array of [file, label] tuples" : s === "ListStringSerializable" ? "array of strings" : s === "FileSerializable" ? "array of files or single file" : e?.description; +} +function Ee(e, s) { + switch (e.msg) { + case "send_data": + return { type: "data" }; + case "send_hash": + return { type: "hash" }; + case "queue_full": + return { + type: "update", + status: { + queue: !0, + message: Ae, + stage: "error", + code: e.code, + success: e.success + } + }; + case "heartbeat": + return { + type: "heartbeat" + }; + case "unexpected_error": + return { + type: "unexpected_error", + status: { + queue: !0, + message: e.message, + session_not_found: e.session_not_found, + stage: "error", + success: !1 + } + }; + case "broken_connection": + return { + type: "broken_connection", + status: { + queue: !0, + message: e.message, + stage: "error", + success: !1 + } + }; + case "estimation": + return { + type: "update", + status: { + queue: !0, + stage: s || "pending", + code: e.code, + size: e.queue_size, + position: e.rank, + eta: e.rank_eta, + success: e.success + } + }; + case "progress": + return { + type: "update", + status: { + queue: !0, + stage: "pending", + code: e.code, + progress_data: e.progress_data, + success: e.success + } + }; + case "log": + return { type: "log", data: e }; + case "process_generating": + return { + type: "generating", + status: { + queue: !0, + message: e.success ? null : e.output.error, + stage: e.success ? "generating" : "error", + code: e.code, + progress_data: e.progress_data, + eta: e.average_duration, + changed_state_ids: e.success ? e.output.changed_state_ids : void 0 + }, + data: e.success ? e.output : null + }; + case "process_streaming": + return { + type: "streaming", + status: { + queue: !0, + message: e.output.error, + stage: "streaming", + time_limit: e.time_limit, + code: e.code, + progress_data: e.progress_data, + eta: e.eta + }, + data: e.output + }; + case "process_completed": + return "error" in e.output ? { + type: "update", + status: { + queue: !0, + title: e.output.title ?? "Error", + message: e.output.error ?? "An error occurred", + visible: e.output.visible, + duration: e.output.duration, + stage: "error", + code: e.code, + success: e.success + } + } : { + type: "complete", + status: { + queue: !0, + message: e.success ? void 0 : e.output.error, + stage: e.success ? "complete" : "error", + code: e.code, + progress_data: e.progress_data, + changed_state_ids: e.success ? e.output.changed_state_ids : void 0 + }, + data: e.success ? e.output : null + }; + case "process_starts": + return { + type: "update", + status: { + queue: !0, + stage: "pending", + code: e.code, + size: e.rank, + position: 0, + success: e.success, + eta: e.eta + }, + original_msg: "process_starts" + }; + } + return { type: "none", status: { stage: "error", queue: !0 } }; +} +const wt = (e = [], s) => { + const t = s ? s.parameters : []; + if (Array.isArray(e)) + return s && t.length > 0 && e.length > t.length && console.warn("Too many arguments provided for the endpoint."), e; + const n = [], o = Object.keys(e); + return t.forEach((r, i) => { + if (e.hasOwnProperty(r.parameter_name)) + n[i] = e[r.parameter_name]; + else if (r.parameter_has_default) + n[i] = r.parameter_default; + else + throw new Error( + `No value provided for required parameter: ${r.parameter_name}` + ); + }), o.forEach((r) => { + if (!t.some((i) => i.parameter_name === r)) + throw new Error( + `Parameter \`${r}\` is not a valid keyword argument. Please refer to the API for usage.` + ); + }), n.forEach((r, i) => { + if (r === void 0 && !t[i].parameter_has_default) + throw new Error( + `No value provided for required parameter: ${t[i].parameter_name}` + ); + }), n; +}; +async function yt() { + if (this.api_info) return this.api_info; + const { token: e } = this.options, { config: s } = this, t = { "Content-Type": "application/json" }; + if (e && (t.Authorization = `Bearer ${e}`), !!s) + try { + let n, o; + if (typeof window < "u" && window.gradio_api_info) + o = window.gradio_api_info; + else { + const r = ce(s.root, this.api_prefix, Ve); + if (n = await this.fetch(r, { + headers: t, + credentials: "include" + }), !n.ok) + throw new Error(L); + o = await n.json(); + } + return "api" in o && (o = o.api), o.named_endpoints["/predict"] && !o.unnamed_endpoints[0] && (o.unnamed_endpoints[0] = o.named_endpoints["/predict"]), _t(o, s, this.api_map); + } catch (n) { + throw new Error("Could not get API info. " + n.message); + } +} +async function bt(e, s, t) { + const n = {}; + this?.options?.token && (n.Authorization = `Bearer ${this.options.token}`); + const o = 1e3, r = []; + let i; + for (let a = 0; a < s.length; a += o) { + const u = s.slice(a, a + o), l = new FormData(); + u.forEach((c) => { + l.append("files", c); + }); + try { + const c = t ? `${e}${this.api_prefix}/${ye}?upload_id=${t}` : `${e}${this.api_prefix}/${ye}`; + i = await this.fetch(c, { + method: "POST", + body: l, + headers: n, + credentials: "include" + }); + } catch (c) { + throw new Error(L + c.message); + } + if (!i.ok) { + const c = await i.text(); + return { error: `HTTP ${i.status}: ${c}` }; + } + const p = await i.json(); + p && r.push(...p); + } + return { files: r }; +} +const vt = { + radix: 1e3, + unit: ["b", "kb", "Mb", "Gb", "Tb", "Pb", "Eb", "Zb", "Yb"] +}, Et = { + radix: 1024, + unit: ["b", "Kib", "Mib", "Gib", "Tib", "Pib", "Eib", "Zib", "Yib"] +}, St = { + radix: 1024, + unit: ["b", "Kb", "Mb", "Gb", "Tb", "Pb", "Eb", "Zb", "Yb"] +}, Se = { + si: vt, + iec: Et, + jedec: St +}; +function $t(e, s = 1, t = "jedec") { + e = Math.abs(e); + const { radix: n, unit: o } = Se[t] || Se.jedec; + let r = 0; + for (; e >= n; ) + e /= n, ++r; + return `${e.toFixed(s)} ${o[r]}`; +} +async function kt(e, s, t, n) { + let o = (Array.isArray(e) ? e : [e]).map( + (i) => i.blob + ); + const r = o.filter( + (i) => i.size > (n ?? 1 / 0) + ); + if (r.length) + throw new Error( + `File(s) exceed the maximum allowed size of ${$t(n || 1 / 0)}: ${r.map((i) => `"${i.name}"`).join(", ")}` + ); + return await Promise.all( + await this.upload_files(s, o, t).then( + async (i) => { + if (i.error) + throw new Error(i.error); + return i.files ? i.files.map((a, u) => new X({ + ...e[u], + path: a, + url: `${s}${this.api_prefix}/file=${a}` + })) : []; + } + ) + ); +} +async function Xt(e, s) { + return e.map( + (t) => new X({ + path: t.name, + orig_name: t.name, + blob: t, + size: t.size, + mime_type: t.type, + is_stream: s + }) + ); +} +class X { + path; + url; + orig_name; + size; + blob; + is_stream; + mime_type; + alt_text; + b64; + meta = { _type: "gradio.FileData" }; + constructor({ + path: s, + url: t, + orig_name: n, + size: o, + blob: r, + is_stream: i, + mime_type: a, + alt_text: u, + b64: l + }) { + this.path = s, this.url = t, this.orig_name = n, this.size = o, this.blob = t ? void 0 : r, this.is_stream = i, this.mime_type = a, this.alt_text = u, this.b64 = l; + } +} +class Le { + type; + command; + meta; + fileData; + constructor(s, t) { + this.type = "command", this.command = s, this.meta = t; + } +} +const Rt = typeof process < "u" && process.versions && process.versions.node; +function $e(e, s, t) { + for (; t.length > 1; ) { + const o = t.shift(); + if (typeof o == "string" || typeof o == "number") + e = e[o]; + else + throw new Error("Invalid key type"); + } + const n = t.shift(); + if (typeof n == "string" || typeof n == "number") + e[n] = s; + else + throw new Error("Invalid key type"); +} +async function ue(e, s = void 0, t = [], n = !1, o = void 0) { + if (Array.isArray(e)) { + let r = []; + return await Promise.all( + e.map(async (i, a) => { + let u = t.slice(); + u.push(String(a)); + const l = await ue( + e[a], + n ? o?.parameters[a]?.component || void 0 : s, + u, + !1, + o + ); + r = r.concat(l); + }) + ), r; + } else { + if (globalThis.Buffer && e instanceof globalThis.Buffer || e instanceof Blob) + return [ + { + path: t, + blob: new Blob([e]), + type: s + } + ]; + if (typeof e == "object" && e !== null) { + let r = []; + for (const i of Object.keys(e)) { + const a = [...t, i], u = e[i]; + r = r.concat( + await ue( + u, + void 0, + a, + !1, + o + ) + ); + } + return r; + } + } + return []; +} +function Ot(e, s) { + let t = s?.dependencies?.find((n) => n.id == e)?.queue; + return t != null ? !t : !s.enable_queue; +} +function Tt(e, s) { + return new Promise((t, n) => { + const o = new MessageChannel(); + o.port1.onmessage = (({ data: r }) => { + o.port1.close(), t(r); + }), window.parent.postMessage(e, s, [o.port2]); + }); +} +function es(e) { + if (typeof e == "string") { + if (e.startsWith("http://") || e.startsWith("https://")) + return { + path: e, + url: e, + orig_name: e.split("/").pop() ?? "unknown", + meta: { _type: "gradio.FileData" } + }; + if (Rt) + return new Le("upload_file", { + path: e, + name: e, + orig_path: e + }); + } else { + if (typeof File < "u" && e instanceof File) + return new Blob([e]); + if (e instanceof Buffer) + return new Blob([e]); + if (e instanceof Blob) + return e; + } + throw new Error( + "Invalid input: must be a URL, File, Blob, or Buffer object." + ); +} +function Y(e, s, t, n, o = !1) { + if (n === "input" && !o) + throw new Error("Invalid code path. Cannot skip state inputs for input."); + if (n === "output" && o) + return e; + let r = [], i = 0; + const a = n === "input" ? s.inputs : s.outputs; + for (let u = 0; u < a.length; u++) { + const l = a[u]; + if (t.find((c) => c.id === l)?.type === "state") { + if (o) + if (e.length === a.length) { + const c = e[i]; + r.push(c), i++; + } else + r.push(null); + else { + i++; + continue; + } + continue; + } else { + const c = e[i]; + r.push(c), i++; + } + } + return r; +} +async function At(e, s, t) { + const n = this; + await xt(n, s); + const o = await ue( + s, + void 0, + [], + !0, + t + ); + return (await Promise.all( + o.map(async ({ path: i, blob: a, type: u }) => { + if (!a) return { path: i, type: u }; + const l = await n.upload_files(e, [a]), p = l.files && l.files[0]; + return { + path: i, + file_url: p, + type: u, + name: typeof File < "u" && a instanceof File ? a?.name : void 0 + }; + }) + )).forEach(({ path: i, file_url: a, type: u, name: l }) => { + if (u === "Gallery") + $e(s, a, i); + else if (a) { + const p = new X({ path: a, orig_name: l }); + $e(s, p, i); + } + }), s; +} +async function xt(e, s) { + if (!(e.config?.root || e.config?.root_url)) + throw new Error(Ce); + await Pe(e, s); +} +async function Pe(e, s, t = []) { + for (const n in s) + s[n] instanceof Le ? await Ct(e, s, n) : typeof s[n] == "object" && s[n] !== null && await Pe(e, s[n], [...t, n]); +} +async function Ct(e, s, t) { + let n = s[t]; + const o = e.config?.root || e.config?.root_url; + if (!o) + throw new Error(Ce); + try { + let r, i; + if (typeof process < "u" && process.versions && process.versions.node) { + const p = await import("./__vite-browser-external-DYxpcVy9.js"); + i = (await import("./__vite-browser-external-DYxpcVy9.js")).resolve(process.cwd(), n.meta.path), r = await p.readFile(i); + } else + throw new Error(ut); + const a = new Blob([r], { + type: "application/octet-stream" + }), u = await e.upload_files(o, [a]), l = u.files && u.files[0]; + if (l) { + const p = new X({ + path: l, + orig_name: n.meta.name || "" + }); + s[t] = p; + } + } catch (r) { + console.error(lt, r); + } +} +async function Dt(e, s, t) { + const n = { "Content-Type": "application/json" }; + this.options.token && (n.Authorization = `Bearer ${this.options.token}`); + try { + var o = await this.fetch(e, { + method: "POST", + body: JSON.stringify(s), + headers: { ...n, ...t }, + credentials: "include" + }); + } catch { + return [{ error: L }, 500]; + } + let r, i; + try { + r = await o.json(), i = o.status; + } catch (a) { + r = { error: `Could not parse server response: ${a}` }, i = 500; + } + return [r, i]; +} +async function Nt(e, s = {}) { + let t = !1, n = !1; + if (!this.config) + throw new Error("Could not resolve app config"); + if (typeof e == "number") + this.config.dependencies.find((o) => o.id == e); + else { + const o = e.replace(/^\//, ""); + this.config.dependencies.find( + (r) => r.id == this.api_map[o] + ); + } + return new Promise(async (o, r) => { + const i = this.submit(e, s, null, null, !0); + let a; + for await (const u of i) + u.type === "data" && (n && o(a), t = !0, a = u), u.type === "status" && (u.stage === "error" && r(u), u.stage === "complete" && (n = !0, t && o(a))); + }); +} +async function H(e, s, t) { + let n = s === "subdomain" ? `https://huggingface.co/api/spaces/by-subdomain/${e}` : `https://huggingface.co/api/spaces/${e}`, o, r; + try { + if (o = await fetch(n), r = o.status, r !== 200) + throw new Error(); + o = await o.json(); + } catch { + t({ + status: "error", + load_status: "error", + message: ot, + detail: "NOT_FOUND" + }); + return; + } + if (!o || r !== 200) return; + const { + runtime: { stage: i }, + id: a + } = o; + switch (i) { + case "STOPPED": + case "SLEEPING": + t({ + status: "sleeping", + load_status: "pending", + message: "Space is asleep. Waking it up...", + detail: i + }), setTimeout(() => { + H(e, s, t); + }, 1e3); + break; + case "PAUSED": + t({ + status: "paused", + load_status: "error", + message: "This space has been paused by the author. If you would like to try this demo, consider duplicating the space.", + detail: i, + discussions_enabled: await ke(a) + }); + break; + case "RUNNING": + case "RUNNING_BUILDING": + t({ + status: "running", + load_status: "complete", + message: "Space is running.", + detail: i + }); + break; + case "BUILDING": + t({ + status: "building", + load_status: "pending", + message: "Space is building...", + detail: i + }), setTimeout(() => { + H(e, s, t); + }, 1e3); + break; + case "APP_STARTING": + t({ + status: "starting", + load_status: "pending", + message: "Space is starting...", + detail: i + }), setTimeout(() => { + H(e, s, t); + }, 1e3); + break; + default: + t({ + status: "space_error", + load_status: "error", + message: "This space is experiencing an issue.", + detail: i, + discussions_enabled: await ke(a) + }); + break; + } +} +const je = async (e, s) => { + let t = 0; + const n = 12, o = 5e3; + return new Promise((r) => { + H( + e, + pe.test(e) ? "space_name" : "subdomain", + (i) => { + s(i), i.status === "running" || i.status === "error" || i.status === "paused" || i.status === "space_error" ? r() : (i.status === "sleeping" || i.status === "building") && (t < n ? (t++, setTimeout(() => { + je(e, s).then(r); + }, o)) : r()); + } + ); + }); +}, Lt = /^(?=[^]*\b[dD]iscussions{0,1}\b)(?=[^]*\b[dD]isabled\b)[^]*$/; +async function ke(e) { + try { + const s = await fetch( + `https://huggingface.co/api/spaces/${e}/discussions`, + { + method: "HEAD" + } + ), t = s.headers.get("x-error-message"); + return !(!s.ok || t && Lt.test(t)); + } catch { + return !1; + } +} +async function Pt(e, s) { + const t = {}; + s && (t.Authorization = `Bearer ${s}`); + try { + const n = await fetch( + `https://huggingface.co/api/spaces/${e}/${Ye}`, + { headers: t } + ); + if (n.status !== 200) + throw new Error("Space hardware could not be obtained."); + const { hardware: o } = await n.json(); + return o.current; + } catch (n) { + throw new Error(n.message); + } +} +async function jt(e, s, t) { + const n = {}; + t && (n.Authorization = `Bearer ${t}`); + const o = { + seconds: s + }; + try { + const r = await fetch( + `https://huggingface.co/api/spaces/${e}/${Qe}`, + { + method: "POST", + headers: { "Content-Type": "application/json", ...n }, + body: JSON.stringify(o) + } + ); + if (r.status !== 200) + throw new Error( + "Could not set sleep timeout on duplicated Space. Please visit *ADD HF LINK TO SETTINGS* to set a timeout manually to reduce billing charges." + ); + return await r.json(); + } catch (r) { + throw new Error(r.message); + } +} +const Re = [ + "cpu-basic", + "cpu-upgrade", + "cpu-xl", + "t4-small", + "t4-medium", + "a10g-small", + "a10g-large", + "a10g-largex2", + "a10g-largex4", + "a100-large", + "zero-a10g", + "h100", + "h100x8" +]; +async function It(e, s) { + const { token: t, private: n, hardware: o, timeout: r, auth: i } = s; + if (o && !Re.includes(o)) + throw new Error( + `Invalid hardware type provided. Valid types are: ${Re.map((g) => `"${g}"`).join(",")}.` + ); + const { http_protocol: a, host: u } = await de(e, t); + let l = null; + if (i) { + const g = await De( + a, + u, + i, + fetch + ); + g && (l = Ne(g)); + } + const p = { + Authorization: `Bearer ${t}`, + "Content-Type": "application/json", + ...l ? { Cookie: l.join("; ") } : {} + }, c = (await (await fetch("https://huggingface.co/api/whoami-v2", { + headers: p + })).json()).name, m = e.split("/")[1], w = { + repository: `${c}/${m}` + }; + n && (w.private = !0); + let P; + try { + o || (P = await Pt(e, t)); + } catch (g) { + throw Error(le + g.message); + } + const f = o || P || "cpu-basic"; + w.hardware = f; + try { + const g = await fetch( + `https://huggingface.co/api/spaces/${e}/duplicate`, + { + method: "POST", + headers: p, + body: JSON.stringify(w) + } + ); + if (g.status === 409) + try { + return await Q.connect(`${c}/${m}`, s); + } catch (z) { + throw console.error("Failed to connect Client instance:", z), z; + } + else if (g.status !== 200) + throw new Error(g.statusText); + const J = await g.json(); + return await jt(`${c}/${m}`, r || 300, t), await Q.connect( + Ut(J.url), + s + ); + } catch (g) { + throw new Error(g); + } +} +function Ut(e) { + const s = /https:\/\/huggingface.co\/spaces\/([^/]+\/[^/]+)/, t = e.match(s); + if (t) + return t[1]; +} +class qt extends TransformStream { + #e = ""; + /** Constructs a new instance. */ + constructor(s = { allowCR: !1 }) { + super({ + transform: (t, n) => { + for (t = this.#e + t; ; ) { + const o = t.indexOf(` +`), r = s.allowCR ? t.indexOf("\r") : -1; + if (r !== -1 && r !== t.length - 1 && (o === -1 || o - 1 > r)) { + n.enqueue(t.slice(0, r)), t = t.slice(r + 1); + continue; + } + if (o === -1) + break; + const i = t[o - 1] === "\r" ? o - 1 : o; + n.enqueue(t.slice(0, i)), t = t.slice(o + 1); + } + this.#e = t; + }, + flush: (t) => { + if (this.#e === "") + return; + const n = s.allowCR && this.#e.endsWith("\r") ? this.#e.slice(0, -1) : this.#e; + t.enqueue(n); + } + }); + } +} +function zt(e) { + let s = new TextDecoderStream(), t = new qt({ allowCR: !0 }); + return e.pipeThrough(s).pipeThrough(t); +} +function Bt(e) { + let t = /[:]\s*/.exec(e), n = t && t.index; + if (n) + return [ + e.substring(0, n), + e.substring(n + t[0].length) + ]; +} +function Oe(e, s, t) { + e.get(s) || e.set(s, t); +} +async function* Ft(e, s) { + if (!e.body) + return; + let t = zt(e.body), n, o = t.getReader(), r; + for (; ; ) { + if (s && s.aborted) + return o.cancel(); + if (n = await o.read(), n.done) + return; + if (!n.value) { + r && (yield r), r = void 0; + continue; + } + let [i, a] = Bt(n.value) || []; + i === "data" ? (r ||= {}, r[i] = r[i] ? r[i] + ` +` + a : a) : i === "event" ? (r ||= {}, r[i] = a) : i === "id" ? (r ||= {}, r[i] = String(+a) === a ? +a : a) : i === "retry" && (r ||= {}, r[i] = +a || void 0); + } +} +async function Gt(e, s) { + let t = new Request(e, s); + Oe(t.headers, "Accept", "text/event-stream"), Oe(t.headers, "Content-Type", "application/json"); + let n = await fetch(t); + if (!n.ok) + throw n; + return Ft(n, t.signal); +} +async function Mt() { + let { + event_callbacks: e, + unclosed_events: s, + pending_stream_messages: t, + stream_status: n, + config: o, + jwt: r + } = this; + const i = this; + if (!o) + throw new Error("Could not resolve app config"); + n.open = !0; + let a = null, u = new URLSearchParams({ + session_hash: this.session_hash + }).toString(), l = new URL(`${o.root}${this.api_prefix}/${Te}?${u}`); + if (r && l.searchParams.set("__sign", r), a = this.stream(l), !a) { + console.warn("Cannot connect to SSE endpoint: " + l.toString()); + return; + } + a.onmessage = async function(p) { + let c = JSON.parse(p.data); + if (c.msg === "close_stream") { + he(n, i.abort_controller); + return; + } + const m = c.event_id; + if (!m) + await Promise.all( + Object.keys(e).map( + (w) => e[w](c) + ) + ); + else if (e[m] && o) { + c.msg === "process_completed" && ["sse", "sse_v1", "sse_v2", "sse_v2.1", "sse_v3"].includes( + o.protocol + ) && s.delete(m); + let w = e[m]; + typeof window < "u" && typeof document < "u" ? setTimeout(w, 0, c) : w(c); + } else + t[m] || (t[m] = []), t[m].push(c); + }, a.onerror = async function(p) { + console.error(p), await Promise.all( + Object.keys(e).map( + (c) => e[c]({ + msg: "broken_connection", + message: L + }) + ) + ); + }; +} +function he(e, s) { + e && (e.open = !1, s?.abort()); +} +function Ht(e, s, t) { + !e[s] ? (e[s] = [], t.data.forEach((o, r) => { + e[s][r] = o; + })) : t.data.forEach((o, r) => { + let i = Jt(e[s][r], o); + e[s][r] = i, t.data[r] = i; + }); +} +function Jt(e, s) { + return s.forEach(([t, n, o]) => { + e = Wt(e, n, t, o); + }), e; +} +function Wt(e, s, t, n) { + if (s.length === 0) { + if (t === "replace") + return n; + if (t === "append") + return e + n; + throw new Error(`Unsupported action: ${t}`); + } + let o = e; + for (let i = 0; i < s.length - 1; i++) + o = o[s[i]]; + const r = s[s.length - 1]; + switch (t) { + case "replace": + o[r] = n; + break; + case "append": + o[r] += n; + break; + case "add": + Array.isArray(o) ? o.splice(Number(r), 0, n) : o[r] = n; + break; + case "delete": + Array.isArray(o) ? o.splice(Number(r), 1) : delete o[r]; + break; + default: + throw new Error(`Unknown action: ${t}`); + } + return e; +} +function Zt(e, s = {}) { + const t = { + close: () => { + console.warn("Method not implemented."); + }, + onerror: null, + onmessage: null, + onopen: null, + readyState: 0, + url: e.toString(), + withCredentials: !1, + CONNECTING: 0, + OPEN: 1, + CLOSED: 2, + addEventListener: () => { + throw new Error("Method not implemented."); + }, + dispatchEvent: () => { + throw new Error("Method not implemented."); + }, + removeEventListener: () => { + throw new Error("Method not implemented."); + } + }; + return Gt(e, s).then(async (n) => { + t.readyState = t.OPEN; + try { + for await (const o of n) + t.onmessage && t.onmessage(o); + t.readyState = t.CLOSED; + } catch (o) { + t.onerror && t.onerror(o), t.readyState = t.CLOSED; + } + }).catch((n) => { + console.error(n), t.onerror && t.onerror(n), t.readyState = t.CLOSED; + }), t; +} +function Kt(e, s = {}, t, n, o, r) { + try { + let i = function(h) { + (o || ze[h.type]) && p(h); + }, a = function() { + for (Me = !0; G.length > 0; ) + G.shift()({ + value: void 0, + done: !0 + }); + }, u = function(h) { + G.length > 0 ? G.shift()(h) : ie.push(h); + }, l = function(h) { + u(Vt(h)), a(); + }, p = function(h) { + u({ value: h, done: !1 }); + }, c = function() { + return ie.length > 0 ? Promise.resolve(ie.shift()) : new Promise((h) => G.push(h)); + }; + const { token: m } = this.options, { + fetch: w, + app_reference: P, + config: f, + session_hash: g, + api_info: J, + api_map: z, + stream_status: ee, + pending_stream_messages: te, + pending_diff_streams: se, + event_callbacks: ne, + unclosed_events: Ie, + post_data: oe, + options: W, + api_prefix: T + } = this, fe = r || { "x-gradio-user": "api" }, Ue = this; + if (!J) throw new Error("No API found"); + if (!f) throw new Error("Could not resolve app config"); + let { fn_index: d, endpoint_info: _e, dependency: B } = Yt( + J, + e, + z, + f + ), qe = wt(s, _e), U, A = f.protocol ?? "ws"; + if (A === "ws") + throw new Error("WebSocket protocol is not supported in this version"); + let q = "", Qt = () => q; + const _ = typeof e == "number" ? "/predict" : e; + let Z, b = null, x = !1, me = {}, F = typeof window < "u" && typeof document < "u" ? new URLSearchParams(window.location.search).toString() : ""; + const ze = W?.events?.reduce( + (h, R) => (h[R] = !0, h), + {} + ) || {}; + async function Be() { + let h = {}, R = {}; + h = { event_id: b }, R = { event_id: b, session_hash: g, fn_index: d }; + try { + if (!f) + throw new Error("Could not resolve app config"); + "event_id" in R && await w(`${f.root}${T}/${st}`, { + headers: { "Content-Type": "application/json" }, + method: "POST", + body: JSON.stringify(R) + }), await w(`${f.root}${T}/${tt}`, { + headers: { "Content-Type": "application/json" }, + method: "POST", + body: JSON.stringify(h) + }); + } catch { + console.warn( + "The `/reset` endpoint could not be called. Subsequent endpoint results may be unreliable." + ); + } + } + const Fe = async (h) => { + await this._resolve_heartbeat(h); + }; + async function ge(h) { + if (!f) return; + let R = h.render_id; + f.components = [ + ...f.components.filter((v) => v.props.rendered_in !== R), + ...h.components + ], f.dependencies = [ + ...f.dependencies.filter((v) => v.rendered_in !== R), + ...h.dependencies + ]; + const K = f.components.some((v) => v.type === "state"), y = f.dependencies.some( + (v) => v.targets.some((j) => j[1] === "unload") + ); + f.connect_heartbeat = K || y, await Fe(f), i({ + type: "render", + data: h, + endpoint: _, + fn_index: d + }); + } + const Ge = this.handle_blob( + f.root, + qe, + _e + ).then(async (h) => { + if (Z = { + data: Y( + h, + B, + f.components, + "input", + !0 + ) || [], + event_data: t, + fn_index: d, + trigger_id: n + }, Ot(d, f)) + i({ + type: "status", + endpoint: _, + stage: "pending", + queue: !1, + fn_index: d, + time: /* @__PURE__ */ new Date() + }), oe( + `${f.root}${T}/run${_.startsWith("/") ? _ : `/${_}`}${F ? "?" + F : ""}`, + { + ...Z, + session_hash: g + }, + fe + ).then(async ([y, v]) => { + const j = y.data; + if (v == 200) + i({ + type: "data", + endpoint: _, + fn_index: d, + data: Y( + j, + B, + f.components, + "output", + W.with_null_state + ), + time: /* @__PURE__ */ new Date(), + event_data: t, + trigger_id: n + }), y.render_config && await ge(y.render_config), i({ + type: "status", + endpoint: _, + fn_index: d, + stage: "complete", + eta: y.average_duration, + queue: !1, + time: /* @__PURE__ */ new Date() + }); + else { + const O = y?.error === L; + i({ + type: "status", + stage: "error", + endpoint: _, + fn_index: d, + message: y.error, + broken: O, + queue: !1, + time: /* @__PURE__ */ new Date() + }); + } + }).catch((y) => { + i({ + type: "status", + stage: "error", + message: y.message, + endpoint: _, + fn_index: d, + queue: !1, + time: /* @__PURE__ */ new Date() + }); + }); + else if (A == "sse") { + i({ + type: "status", + stage: "pending", + queue: !0, + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date() + }); + var K = new URLSearchParams({ + fn_index: d.toString(), + session_hash: g + }).toString(); + let y = new URL( + `${f.root}${T}/${Te}?${F ? F + "&" : ""}${K}` + ); + if (this.jwt && y.searchParams.set("__sign", this.jwt), U = this.stream(y), !U) + return Promise.reject( + new Error("Cannot connect to SSE endpoint: " + y.toString()) + ); + U.onmessage = async function(v) { + const j = JSON.parse(v.data), { type: O, status: C, data: D } = Ee( + j, + me[d] + ); + if (O === "update" && C && !x) + i({ + type: "status", + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date(), + ...C + }), C.stage === "error" && (U?.close(), a()); + else if (O === "data") { + let [k, I] = await oe( + `${f.root}${T}/queue/data`, + { + ...Z, + session_hash: g, + event_id: b + } + ); + I !== 200 && (i({ + type: "status", + stage: "error", + message: L, + queue: !0, + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date() + }), U?.close(), a()); + } else O === "complete" ? x = C : O === "log" ? i({ + type: "log", + title: D.title, + log: D.log, + level: D.level, + endpoint: _, + duration: D.duration, + visible: D.visible, + fn_index: d + }) : (O === "generating" || O === "streaming") && i({ + type: "status", + time: /* @__PURE__ */ new Date(), + ...C, + stage: C?.stage, + queue: !0, + endpoint: _, + fn_index: d + }); + D && (i({ + type: "data", + time: /* @__PURE__ */ new Date(), + data: Y( + D.data, + B, + f.components, + "output", + W.with_null_state + ), + endpoint: _, + fn_index: d, + event_data: t, + trigger_id: n + }), x && (i({ + type: "status", + time: /* @__PURE__ */ new Date(), + ...x, + stage: C?.stage, + queue: !0, + endpoint: _, + fn_index: d + }), U?.close(), a())); + }; + } else if (A == "sse_v1" || A == "sse_v2" || A == "sse_v2.1" || A == "sse_v3") { + i({ + type: "status", + stage: "pending", + queue: !0, + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date() + }); + let y = ""; + typeof window < "u" && typeof document < "u" && (y = window?.location?.hostname); + const j = y.includes(".dev.") ? `https://moon-${y.split(".")[1]}.dev.spaces.huggingface.tech` : "https://huggingface.co"; + return (typeof window < "u" && typeof document < "u" && window.parent != window && window.supports_zerogpu_headers ? Tt("zerogpu-headers", j) : Promise.resolve(null)).then((k) => { + const I = { ...fe, ...k || {} }; + return oe( + `${f.root}${T}/${Ze}?${F}`, + { + ...Z, + session_hash: g + }, + I + ); + }).then(async ([k, I]) => { + if (k.event_id && (q = k.event_id), I === 503) + i({ + type: "status", + stage: "error", + message: Ae, + queue: !0, + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date(), + visible: !0 + }); + else if (I === 422) + i({ + type: "status", + stage: "error", + message: k.detail, + queue: !0, + endpoint: _, + fn_index: d, + code: "validation_error", + time: /* @__PURE__ */ new Date(), + visible: !0 + }), a(); + else if (I !== 200) { + const M = k?.error === L; + i({ + type: "status", + stage: "error", + broken: M, + message: M ? L : k.detail || k.error, + queue: !0, + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date(), + visible: !0 + }); + } else { + b = k.event_id, q = b; + let M = async function(re) { + try { + const { type: E, status: S, data: $, original_msg: He } = Ee( + re, + me[d] + ); + if (E == "heartbeat") + return; + if (E === "update" && S && !x) + i({ + type: "status", + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date(), + original_msg: He, + ...S + }); + else if (E === "complete") + x = S; + else if (E == "unexpected_error" || E == "broken_connection") { + console.error("Unexpected error", S?.message); + const Je = E === "broken_connection"; + i({ + type: "status", + stage: "error", + message: S?.message || "An Unexpected Error Occurred!", + queue: !0, + endpoint: _, + broken: Je, + session_not_found: S?.session_not_found, + fn_index: d, + time: /* @__PURE__ */ new Date() + }); + } else if (E === "log") { + i({ + type: "log", + title: $.title, + log: $.log, + level: $.level, + endpoint: _, + duration: $.duration, + visible: $.visible, + fn_index: d + }); + return; + } else (E === "generating" || E === "streaming") && (i({ + type: "status", + time: /* @__PURE__ */ new Date(), + ...S, + stage: S?.stage, + queue: !0, + endpoint: _, + fn_index: d + }), $ && B.connection !== "stream" && ["sse_v2", "sse_v2.1", "sse_v3"].includes(A) && Ht(se, b, $)); + $ && (i({ + type: "data", + time: /* @__PURE__ */ new Date(), + data: Y( + $.data, + B, + f.components, + "output", + W.with_null_state + ), + endpoint: _, + fn_index: d + }), $.render_config && await ge($.render_config), x && (i({ + type: "status", + time: /* @__PURE__ */ new Date(), + ...x, + stage: S?.stage, + queue: !0, + endpoint: _, + fn_index: d + }), a())), (S?.stage === "complete" || S?.stage === "error") && (ne[b] && delete ne[b], b in se && delete se[b]); + } catch (E) { + console.error("Unexpected client exception", E), i({ + type: "status", + stage: "error", + message: "An Unexpected Error Occurred!", + queue: !0, + endpoint: _, + fn_index: d, + time: /* @__PURE__ */ new Date() + }), ["sse_v2", "sse_v2.1", "sse_v3"].includes(A) && (he(ee, Ue.abort_controller), ee.open = !1, a()); + } + }; + b in te && (te[b].forEach((re) => M(re)), delete te[b]), ne[b] = M, Ie.add(b), ee.open || await this.open_stream(); + } + }); + } + }); + let Me = !1; + const ie = [], G = [], we = { + [Symbol.asyncIterator]: () => we, + next: c, + throw: async (h) => (l(h), c()), + return: async () => (a(), { value: void 0, done: !0 }), + cancel: Be, + send_chunk: (h) => { + this.post_data(`${f.root}${T}/stream/${q}`, { + ...h, + session_hash: this.session_hash + }); + }, + close_stream: () => { + this.post_data( + `${f.root}${T}/stream/${q}/close`, + {} + ), a(); + }, + event_id: () => q, + wait_for_id: async () => (await Ge, b) + }; + return we; + } catch (i) { + throw console.error("Submit function encountered an error:", i), i; + } +} +function Vt(e) { + return { + then: (s, t) => t(e) + }; +} +function Yt(e, s, t, n) { + let o, r, i; + if (typeof s == "number") + o = s, r = e.unnamed_endpoints[o], i = n.dependencies.find((a) => a.id == s); + else { + const a = s.replace(/^\//, ""); + o = t[a], r = e.named_endpoints[s.trim()], i = n.dependencies.find( + (u) => u.id == t[a] + ); + } + if (typeof o != "number") + throw new Error( + "There is no endpoint matching that name of fn_index matching that number." + ); + return { fn_index: o, endpoint_info: r, dependency: i }; +} +class Q { + app_reference; + options; + deep_link = null; + config; + api_prefix = ""; + api_info; + api_map = {}; + session_hash = Math.random().toString(36).substring(2); + jwt = !1; + last_status = {}; + cookies = null; + // streaming + stream_status = { open: !1 }; + closed = !1; + pending_stream_messages = {}; + pending_diff_streams = {}; + event_callbacks = {}; + unclosed_events = /* @__PURE__ */ new Set(); + heartbeat_event = null; + abort_controller = null; + stream_instance = null; + current_payload; + get_url_config(s = null) { + if (!this.config) + throw new Error(N); + s === null && (s = window.location.href); + const t = (i) => i.replace(/^\/+|\/+$/g, ""); + let n = t(new URL(this.config.root).pathname), o = t(new URL(s).pathname), r; + return o.startsWith(n) ? r = t(o.substring(n.length)) : r = "", this.get_page_config(r); + } + get_page_config(s) { + if (!this.config) + throw new Error(N); + let t = this.config; + return s in t.page || (s = ""), { + ...t, + current_page: s, + layout: t.page[s].layout, + components: t.components.filter( + (n) => t.page[s].components.includes(n.id) + ), + dependencies: this.config.dependencies.filter( + (n) => t.page[s].dependencies.includes(n.id) + ) + }; + } + fetch(s, t) { + const n = new Headers(t?.headers || {}); + return this && this.cookies && n.append("Cookie", this.cookies), this && this.options.headers && new Headers(this.options.headers).forEach((r, i) => { + n.append(i, r); + }), fetch(s, { ...t, headers: n }); + } + stream(s) { + const t = new Headers(); + return this && this.cookies && t.append("Cookie", this.cookies), this && this.options.headers && new Headers(this.options.headers).forEach((o, r) => { + t.append(r, o); + }), this && this.options.token && t.append("Authorization", `Bearer ${this.options.token}`), this.abort_controller = new AbortController(), this.stream_instance = Zt(s.toString(), { + credentials: "include", + headers: t, + signal: this.abort_controller.signal + }), this.stream_instance; + } + view_api; + upload_files; + upload; + handle_blob; + post_data; + submit; + predict; + open_stream; + resolve_config; + resolve_cookies; + constructor(s, t = { events: ["data"] }) { + this.app_reference = s, this.deep_link = t.query_params?.deep_link || null, t.events || (t.events = ["data"]), this.options = t, this.current_payload = {}, t.cookies && (this.cookies = t.cookies), this.view_api = yt.bind(this), this.upload_files = bt.bind(this), this.handle_blob = At.bind(this), this.post_data = Dt.bind(this), this.submit = Kt.bind(this), this.predict = Nt.bind(this), this.open_stream = Mt.bind(this), this.resolve_config = dt.bind(this), this.resolve_cookies = ht.bind(this), this.upload = kt.bind(this), this.fetch = this.fetch.bind(this), this.handle_space_success = this.handle_space_success.bind(this), this.stream = this.stream.bind(this); + } + async init() { + this.options.auth && await this.resolve_cookies(), await this._resolve_config().then( + ({ config: s }) => this._resolve_heartbeat(s) + ), this.api_info = await this.view_api(), this.api_map = pt(this.config?.dependencies || []); + } + async _resolve_heartbeat(s) { + if (s && (this.config = s, this.api_prefix = s.api_prefix || "", this.config && this.config.connect_heartbeat && this.config.space_id && this.options.token && (this.jwt = await be( + this.config.space_id, + this.options.token, + this.cookies + ))), s.space_id && this.options.token && (this.jwt = await be(s.space_id, this.options.token)), this.config && this.config.connect_heartbeat) { + const t = new URL( + `${this.config.root}${this.api_prefix}/${Xe}/${this.session_hash}` + ); + this.jwt && t.searchParams.set("__sign", this.jwt), this.heartbeat_event || (this.heartbeat_event = this.stream(t)); + } + } + static async connect(s, t = { + events: ["data"] + }) { + const n = new this(s, t); + return t.session_hash && (n.session_hash = t.session_hash), await n.init(), n; + } + async reconnect() { + const s = new URL( + `${this.config.root}${this.api_prefix}/${nt}` + ); + let t; + try { + const n = await this.fetch(s); + if (!n.ok) + throw new Error(); + t = (await n.json()).app_id; + } catch { + return "broken"; + } + return t !== this.config.app_id ? "changed" : "connected"; + } + close() { + this.closed = !0, he(this.stream_status, this.abort_controller); + } + set_current_payload(s) { + this.current_payload = s; + } + static async duplicate(s, t = { + events: ["data"] + }) { + return It(s, t); + } + async _resolve_config() { + const { http_protocol: s, host: t, space_id: n } = await de( + this.app_reference, + this.options.token + ), { status_callback: o } = this.options; + n && o && await je(n, o); + let r; + try { + let i = `${s}//${t}`; + if (r = await this.resolve_config(i), !r) + throw new Error(N); + return this.config_success(r); + } catch (i) { + if (n && o) + H( + n, + pe.test(n) ? "space_name" : "subdomain", + this.handle_space_success + ); + else + throw o && o({ + status: "error", + message: "Could not load this space.", + load_status: "error", + detail: "NOT_FOUND" + }), Error(i); + } + } + async config_success(s) { + if (this.config = s, this.api_prefix = s.api_prefix || "", this.config.auth_required) + return this.prepare_return_obj(); + try { + this.api_info = await this.view_api(); + } catch (t) { + console.error(it + t.message); + } + return this.prepare_return_obj(); + } + async handle_space_success(s) { + if (!this) + throw new Error(N); + const { status_callback: t } = this.options; + if (t && t(s), s.status === "running") + try { + if (this.config = await this._resolve_config(), this.api_prefix = this?.config?.api_prefix || "", !this.config) + throw new Error(N); + return await this.config_success(this.config); + } catch (n) { + throw t && t({ + status: "error", + message: "Could not load this space.", + load_status: "error", + detail: "NOT_FOUND" + }), n; + } + } + async component_server(s, t, n) { + if (!this.config) + throw new Error(N); + const o = {}, { token: r } = this.options, { session_hash: i } = this; + r && (o.Authorization = `Bearer ${this.options.token}`); + let a, u = this.config.components.find( + (p) => p.id === s + ); + u?.props?.root_url ? a = u.props.root_url : a = this.config.root; + let l; + if (typeof n == "object" && n !== null && "binary" in n) { + const p = n; + l = new FormData(); + for (const c in p.data) + c !== "binary" && l.append(c, p.data[c]); + l.set("component_id", s.toString()), l.set("fn_name", t), l.set("session_hash", i); + } else + l = JSON.stringify({ + data: n, + component_id: s, + fn_name: t, + session_hash: i + }), o["Content-Type"] = "application/json"; + r && (o.Authorization = `Bearer ${r}`); + try { + const p = await this.fetch( + `${a}${this.api_prefix}/${et}/`, + { + method: "POST", + body: l, + headers: o, + credentials: "include" + } + ); + if (!p.ok) + throw new Error( + "Could not connect to component server: " + p.statusText + ); + return await p.json(); + } catch (p) { + console.warn(p); + } + } + set_cookies(s) { + this.cookies = Ne(s).join("; "); + } + prepare_return_obj() { + return { + config: this.config, + predict: this.predict, + submit: this.submit, + view_api: this.view_api, + component_server: this.component_server + }; + } +} +async function ts(e, s = { + events: ["data"] +}) { + return await Q.connect(e, s); +} +async function ss(e, s) { + return await Q.duplicate(e, s); +} +export { + Q as Client, + X as FileData, + ct as MISSING_CREDENTIALS_MSG, + ts as client, + ss as duplicate, + es as handle_file, + Nt as predict, + Xt as prepare_files, + Kt as submit, + kt as upload, + bt as upload_files +}; diff --git a/hf_space/vendor/gradio-client-LICENSE b/hf_space/vendor/gradio-client-LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/hf_space/vendor/gradio-client-LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/hf_space/vendor/gradio-client-source.json b/hf_space/vendor/gradio-client-source.json new file mode 100644 index 0000000000000000000000000000000000000000..b529473e4a8cfe4349d84ba8633244a9d5ce81b8 --- /dev/null +++ b/hf_space/vendor/gradio-client-source.json @@ -0,0 +1,9 @@ +{ + "package": "@gradio/client", + "version": "2.1.0", + "matching_gradio": "6.10.0", + "source": "https://cdn.jsdelivr.net/npm/@gradio/client@2.1.0/dist/browser.js", + "sha256": "2a88bfa88b6c3ac52241287a36bf0a6325996d815abe0218a4be037f5067b3ea", + "license_file": "gradio-client-LICENSE", + "modified": false +} diff --git a/hf_space/viewport.js b/hf_space/viewport.js index c5a0aac312b33c6de266db3ad2374120d0db1087..5e2fb04895277d2a7a6e4bf65e4e5ead0f174e5f 100644 --- a/hf_space/viewport.js +++ b/hf_space/viewport.js @@ -36,8 +36,8 @@ // Leave room for HF's wrapper. There is no fixed minimum desktop height. const height = Math.max(1, Math.floor(available) - (embedded ? 16 : 0)); document.documentElement.style.setProperty('--st-available-height', height + 'px'); - // Measure the header's actual position, the resource notice (including - // its margins), and Gradio's footer. Never budget using iframe height, + // Measure the header's position, plan notice, and Gradio's footer. + // Never budget using iframe height, // which grows with its content and would create a resize feedback loop. const top = root.getBoundingClientRect().top + window.scrollY; const tail = outerHeight(resourceFooter()) + gradioFooters().reduce((sum, node) => sum + outerHeight(node), 0); diff --git a/hf_space/zerogpu_client.js b/hf_space/zerogpu_client.js new file mode 100644 index 0000000000000000000000000000000000000000..a66a6ecc868047a3aab1c446d3187e5c0c2c3371 --- /dev/null +++ b/hf_space/zerogpu_client.js @@ -0,0 +1,295 @@ +/* Request-scoped ZeroGPU dispatch. HF quota headers stay in memory for one call. */ +(() => { + const fallback = 'The GPU connection was interrupted. Refresh the page to reconnect your saved request.'; + const identityError = () => new Error('Hugging Face could not verify your personal ZeroGPU quota.'); + const quotaToken = headers => headers.get('x-zerogpu-token')?.trim() || headers.get('x-ip-token')?.trim(); + const modelsReady = state => state?.native_build?.stage === 'compiled' && state?.models?.stage === 'ready_released'; + const workspaceReady = state => state?.gpu_mode === 'zerogpu' ? state.workspace?.stage === 'ready' : + modelsReady(state) && (!state.workspace || state.workspace.stage === 'ready'); + const foregroundOperation = record => record?.gpu_operation_kind === 'foreground' || !!record?.foreground; + const preparationState = (state, record) => (foregroundOperation(record) ? state?.foreground : + state?.native_build && state.native_build.stage !== 'compiled' ? state.native_build : state?.models) || {}; + const operationReady = (state, record) => workspaceReady(state) && (foregroundOperation(record) ? + state?.foreground?.stage === 'ready_released' : modelsReady(state)); + const preparationError = (state, record) => { + const failed = value => /^(failed|blocked)/.test(value?.stage || ''); + return foregroundOperation(record) ? failed(state?.foreground) ? 'Background-removal model preparation failed' : null : + failed(state?.native_build) || failed(state?.models) ? 'Model preparation failed' : null; + }; + function formatProgress(progress = {}) { + const count = value => Number.isFinite(Number(value)) ? Math.max(0, Number(value)) : 0; + const total = count(progress.total), completed = Math.min(count(progress.completed), total); + const percent = total ? Math.round(completed / total * 1000) / 10 : 0; + let quantity = ''; + if (total) { + const divisor = total >= 1e9 ? 1e9 : 1e6; + quantity = progress.unit === 'bytes' ? `${(completed / divisor).toFixed(1)} / ${(total / divisor).toFixed(1)} ${divisor === 1e9 ? 'GB' : 'MB'}` : + `${completed} / ${total} ${progress.unit === 'models' ? 'models' : 'steps'}`; + } + return {percent, text: `${percent}%${quantity ? ' · ' + quantity : ''}`}; + } + function safeMessage(error) { + const value = typeof error === 'string' ? error : error?.message; + if (typeof value !== 'string') return fallback; + // Do not display arbitrary backend exception text, paths, or credentials. + // A missing/invalid identity is not depleted quota. In particular, HF may + // return an IP-fallback token even while the account has its full allowance. + if (/could not verify|cannot verify|unable to verify|verification|sign.in|login|log in|authenticated|session.*expir/i.test(value)) + return 'Hugging Face could not verify your account for ZeroGPU. Sign in on Hugging Face, then refresh this Space. This does not mean your quota is exhausted.'; + if (/background-removal model preparation failed/i.test(value)) + return 'The background-removal model could not be prepared. Please retry the task.'; + if (/model preparation failed/i.test(value)) + return 'Models could not be prepared for this task. Please retry the task.'; + if (/quota[^.\n]*(?:exhausted|exceeded|insufficient)|(?:exhausted|exceeded|insufficient|not enough)[^.\n]*quota|(?:daily GPU|GPU usage)[^.\n]*(?:exceeded|exhausted|limit reached)/i.test(value)) + return 'Your available Hugging Face GPU quota is insufficient for this task. Check your personal quota before trying again.'; + if (/duration|timeout|timed out|time limit/i.test(value)) + return 'This task exceeded its GPU time limit. Please try again.'; + if (/busy|queue.*full|no.*GPU.*available/i.test(value)) + return 'ZeroGPU is busy. Please try again shortly.'; + if (/quota/i.test(value)) + return 'Hugging Face could not process this GPU quota request. Refresh this Space and try again; your quota may still be available.'; + return fallback; + } + + function waitForVisible(signal) { + signal.throwIfAborted(); + if (!document.hidden) return Promise.resolve(); + return new Promise((resolve, reject) => { + const finish = error => { + document.removeEventListener('visibilitychange', visible); + signal.removeEventListener('abort', cancelled); + error ? reject(error) : resolve(); + }; + const visible = () => { if (!document.hidden) finish(); }; + const cancelled = () => finish(signal.reason); + document.addEventListener('visibilitychange', visible); + signal.addEventListener('abort', cancelled, {once: true}); + }); + } + + async function requestQuotaHeaders(signal) { + if (typeof window === 'undefined' || window.parent === window) throw identityError(); + const origin = 'https://huggingface.co'; + // An app OAuth cookie does not establish the embedding HF page's identity. + // Ask that page explicitly, only while this document is visible. + while (true) { + await waitForVisible(signal); + signal.throwIfAborted(); + await new Promise((resolve, reject) => { + const finish = error => { + clearTimeout(timer); window.removeEventListener('message', supported); + signal.removeEventListener('abort', cancelled); + error ? reject(error) : resolve(); + }; + const supported = event => { + if (event.source === window.parent && event.origin === origin && event.data === 'supports-zerogpu-headers') finish(); + }; + const cancelled = () => finish(signal.reason); + const timer = setTimeout(() => finish(identityError()), 10000); + window.addEventListener('message', supported); + signal.addEventListener('abort', cancelled, {once: true}); + try { window.parent.postMessage('supports-zerogpu-headers', origin); } + catch { finish(identityError()); } + }); + if (document.hidden) continue; + signal.throwIfAborted(); + const headers = await new Promise((resolve, reject) => { + const channel = new MessageChannel(); + const finish = (error, value) => { + clearTimeout(timer); channel.port1.close(); channel.port2.close(); + signal.removeEventListener('abort', cancelled); + error ? reject(error) : resolve(value); + }; + const cancelled = () => finish(signal.reason); + const timer = setTimeout(() => finish(identityError()), 10000); + channel.port1.onmessage = ({data}) => { + try { + const values = new Headers(data || {}); + // HF returns X-ZeroGPU-Token/UUID here; its proxy derives the + // x-ip-token seen by Python. Older parents may return x-ip-token. + if (!quotaToken(values)) return document.hidden ? finish(null, null) : finish(identityError()); + values.set('x-gradio-user', 'app'); + finish(null, Object.fromEntries(values)); + } catch { finish(identityError()); } + }; + signal.addEventListener('abort', cancelled, {once: true}); + try { window.parent.postMessage('zerogpu-headers', origin, [channel.port2]); } + catch { finish(identityError()); } + }); + // HF can omit the token if visibility changed mid-request. The null + // marker retries even if visibility returns before this continuation. + if (headers && !document.hidden) return headers; + } + } + + function createDispatcher({identity, version, stopped, ready = () => true, prepare = async () => {}, + quotaHeaders = requestQuotaHeaders, onPending, onError, onComplete, + connect = async options => { + const {Client} = await import('/gradio_api/file=hf_space/vendor/gradio-client-2.1.0.js'); + return Client.connect(location.origin, options); + }}) { + const records = new Map(), clients = new Set(); + let epoch = 0, runtimeId = null; + const current = (generation, accountVersion) => + generation === epoch && accountVersion === version() && !stopped() && identity()?.authenticated; + + function valid(record) { + const id = record?.gpu_operation_id; + return record?.gpu_mode === 'zerogpu' && typeof id === 'string' && /^[a-f0-9]{32}$/.test(id) && + identity()?.authenticated && !stopped(); + } + function fail(entry, error) { + if (!current(entry.generation, entry.accountVersion)) return; + entry.phase = 'failed'; entry.failure = safeMessage(error); + onError?.(entry.record, entry.failure); + } + function start(entry) { + if (entry.phase !== 'waiting' || !ready(entry.record) || !current(entry.generation, entry.accountVersion)) return entry.promise; + entry.phase = 'connecting'; + return entry.promise = (async () => { + let connection; + try { + let values = new Headers(await quotaHeaders(entry.controller.signal) || {}); + if (!quotaToken(values)) throw identityError(); + if (!current(entry.generation, entry.accountVersion)) return; + const csrf = identity()?.csrf_token; + if (!csrf) throw new Error('Sign in again'); + // Each operation owns its client, including its in-memory headers. + // Do not put x-gradio-user in options: Client.fetch appends them. + connection = Promise.resolve(connect({events: ['data', 'status'], headers: {'X-CSRF-Token': csrf}})); + clients.add(connection); + const client = await connection; + if (!current(entry.generation, entry.accountVersion)) return; + if (typeof document !== 'undefined' && document.hidden) { + values = new Headers(await quotaHeaders(entry.controller.signal) || {}); + if (!quotaToken(values)) throw identityError(); + if (!current(entry.generation, entry.accountVersion)) return; + } + values.set('x-gradio-user', 'app'); + const headers = Object.fromEntries(values); + const originalFetch = client.fetch; + if (originalFetch) client.fetch = async (url, options) => { + const target = new URL(url, location.origin); + if (target.origin === location.origin && target.pathname.endsWith('/queue/join')) { + // Client 2.1.0 repeats the iframe handshake and merges its result + // over submit's headers. Recheck visibility and obtain fresh + // headers at the actual send, after that second handshake. + try { + const fresh = new Headers(await quotaHeaders(entry.controller.signal) || {}); + if (!quotaToken(fresh)) throw identityError(); + entry.controller.signal.throwIfAborted(); + if (!current(entry.generation, entry.accountVersion)) throw identityError(); + const locked = new Headers(options?.headers); + for (const name of ['x-zerogpu-token', 'x-zerogpu-uuid', 'x-ip-token']) locked.delete(name); + for (const [name, value] of fresh) locked.set(name, value); + locked.set('x-gradio-user', 'app'); + return originalFetch.call(client, url, {...options, headers: locked}); + } catch (error) { + // post_data catches fetch exceptions; retain only in memory + // so its generic transport error keeps the right UI category. + entry.headerError = error; throw error; + } + } + return originalFetch.call(client, url, options); + }; + const submission = client.submit('/run_gpu_operation', {operation_id: entry.record.gpu_operation_id}, + undefined, undefined, undefined, headers); + entry.phase = 'running'; + let result = null; + for await (const event of submission) { + if (!current(entry.generation, entry.accountVersion)) return; + if (event.type === 'status' && event.stage === 'error') throw entry.headerError || event; + if (event.type === 'data') { + result = event.data?.[0] ?? event.data; + if (result?.status === 'queued') throw new Error('ZeroGPU is busy'); + if (result?.status === 'failed' || result?.status === 'error') + throw new Error(result.error || result.message || 'GPU task failed'); + } + } + if (current(entry.generation, entry.accountVersion)) { + entry.phase = 'complete'; onComplete?.(entry.record, result); + } + } catch (error) { fail(entry, error); } + finally { + if (connection && clients.delete(connection)) connection.then(client => client.close?.()).catch(() => {}); + } + })(); + } + function prepareEntry(entry) { + const attempt = {}; + entry.preparation = attempt; + entry.controller.abort(); entry.controller = new AbortController(); + entry.phase = 'preparing'; + const active = () => entry.preparation === attempt && + records.get(entry.record.gpu_operation_id) === entry && current(entry.generation, entry.accountVersion); + entry.promise = (async () => { + try { + // Only explicit user intent reaches this authenticated CPU-only POST. + await prepare(entry.record, entry.controller.signal); + if (!active() || entry.phase === 'failed') return; + entry.phase = 'waiting'; onPending?.(entry.record); + if (ready(entry.record)) await start(entry); + } catch (error) { if (active()) fail(entry, error); } + })(); + return entry.promise; + } + function dispatch(record) { + if (!valid(record)) return null; + const id = record.gpu_operation_id; + if (records.has(id)) return start(records.get(id)); + const entry = {record, generation: epoch, accountVersion: version(), controller: new AbortController(), + phase: 'preparing', failure: '', promise: null}; + records.set(id, entry); + return prepareEntry(entry); + } + + return { + observe(response, {requested = false} = {}) { + for (const record of [response, ...(response?.items || [])]) { + if (!valid(record)) continue; + const entry = records.get(record.gpu_operation_id); + if (entry) { entry.record = record; start(entry); } + else if (requested) dispatch(record); + } + }, + dispatch, + syncRuntime(value) { + if (typeof value !== 'string' || !value || value === runtimeId) return; + const previous = runtimeId; runtimeId = value; + if (previous === null) return; + // A Space restart loses its CPU preparation controllers. Re-arm only + // requests explicitly started in this page and not yet sent to a GPU. + // History reads and already submitted GPU work must never be replayed. + for (const entry of records.values()) { + if (['preparing', 'waiting'].includes(entry.phase) && current(entry.generation, entry.accountVersion)) + void prepareEntry(entry); + } + }, + advance(error) { + for (const entry of records.values()) { + const failure = typeof error === 'function' ? error(entry.record) : error; + if (failure && ['preparing', 'waiting'].includes(entry.phase)) fail(entry, failure); + else start(entry); + } + }, + requested(id) { return records.has(id); }, + pending(predicate = () => true) { + return [...records.values()].some(entry => ['preparing', 'waiting', 'connecting', 'running'].includes(entry.phase) && predicate(entry.record)); + }, + failure(id) { return records.get(id)?.failure || ''; }, + retry(record) { + if (records.get(record?.gpu_operation_id)?.failure) records.delete(record.gpu_operation_id); + return dispatch(record); + }, + reset() { + epoch++; + for (const entry of records.values()) entry.controller.abort(); + records.clear(); + for (const client of clients) Promise.resolve(client).then(value => value.close?.()).catch(() => {}); + clients.clear(); + } + }; + } + globalThis.SupraZeroGPU = Object.freeze({createDispatcher, requestQuotaHeaders, modelsReady, workspaceReady, + foregroundOperation, preparationState, operationReady, preparationError, formatProgress}); +})(); diff --git a/hf_space/zerogpu_jobs.py b/hf_space/zerogpu_jobs.py new file mode 100644 index 0000000000000000000000000000000000000000..6a69be857ae983989ca47cfaf57baffa803a90fb --- /dev/null +++ b/hf_space/zerogpu_jobs.py @@ -0,0 +1,78 @@ +"""Authenticated Gradio entry to durable operations; no stored quota tokens.""" +import base64 +import json +import logging + +import gradio as gr +from fastapi import HTTPException + +from hf_space.history_auth import require_mutation +from hf_space.history_store import HistoryError, NotFound + +LOG = logging.getLogger(__name__) + + +def _has_quota_error(token): + # Match spaces 0.51.3's jwt_payload parser and truthy error check. This + # unverified payload can only reject a known IP-quota fallback; it cannot + # establish identity. Unknown token formats remain HF's responsibility. + try: + _, payload, _ = token.split('.') + value = json.loads(base64.urlsafe_b64decode(payload + '==')) + except (ValueError, TypeError): + return False + return isinstance(value, dict) and bool(value.get('error')) + + +def public_gpu_error(error, fallback): + # Never persist scheduler tracebacks, request headers or arbitrary exceptions. + message = str(error).lower() + if any(value in message for value in ('verify', 'verification', 'authenticate', 'sign in', 'login', 'invalid token', 'expired token')): + return 'Hugging Face could not verify your personal ZeroGPU session. Keep the Space page visible and reconnect your account.' + if any(value in message for value in ('quota exceeded', 'quota exhausted', 'quota is insufficient', 'insufficient quota', 'daily gpu limit')): + return 'Your Hugging Face ZeroGPU quota is insufficient. Wait for it to reset, then try again.' + if 'duration' in message or 'timed out' in message or 'timeout' in message: + return 'The ZeroGPU time limit was reached. Please try again with your available GPU quota.' + if 'zerogpu' in message or 'gpu task aborted' in message: + return 'ZeroGPU is unavailable for this request. Please try again shortly.' + return fallback + + +def register(demo, service): + def run_gpu_operation(operation_id: str, request: gr.Request): + try: + if request is None or request.request is None: + raise HTTPException(401, 'Sign in with Hugging Face before using the GPU.') + identity = require_mutation(request.request) + # HF validates this opaque token when scheduling. OAuth alone does + # not identify the account whose GPU allowance will be consumed. + quota_token = dict(request.headers).get('x-ip-token') + if not quota_token: + raise HTTPException(403, 'Open this Space on Hugging Face and sign in to use your own ZeroGPU quota.') + if _has_quota_error(quota_token): + # Diagnostic booleans only: never log credentials or claims. + LOG.warning('ZeroGPU identity rejected: browser_quota_header=%s, browser_session_header=%s, app_request=%s', + bool(request.headers.get('x-zerogpu-token')), + bool(request.headers.get('x-zerogpu-uuid')), + request.headers.get('x-gradio-user') == 'app') + raise HTTPException(403, 'Hugging Face could not verify your personal ZeroGPU quota. Sign in on Hugging Face, then refresh this Space.') + pending = service.store.get_operation(identity.owner, operation_id) + if not service.ready(pending['operation']['kind']): + raise HTTPException(503, 'Models are loading. Please try again when loading is complete.') + service.runner.run_once(operation_id=operation_id, owner=identity.owner) + current = service.store.get_operation(identity.owner, operation_id)['operation'] + return {'operation_id': operation_id, 'status': current['status'], 'error': current.get('error')} + except HTTPException as error: + raise gr.Error(error.detail) from None + except (NotFound, HistoryError): + raise gr.Error('This task is not available for your account.') from None + + with demo: + operation = gr.Textbox(visible=False) + result = gr.JSON(visible=False) + trigger = gr.Button(visible=False) + trigger.click(run_gpu_operation, operation, result, + api_name='run_gpu_operation', api_visibility='public', + concurrency_limit=None, + show_progress='hidden') + return run_gpu_operation diff --git a/hf_space/zerogpu_runtime.py b/hf_space/zerogpu_runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..901c2eaeab0ef90a452da4620842c6f7492e3f57 --- /dev/null +++ b/hf_space/zerogpu_runtime.py @@ -0,0 +1,149 @@ +"""GPU leases for inference; the web process owns only immutable CPU weights. + +Only primitive paths, seeds and progress identifiers cross the ZeroGPU process +boundary. History stores, locks, callbacks and authentication stay in the web +process. The caller must invoke these from its authenticated Gradio request. +""" +from contextlib import redirect_stderr, redirect_stdout +import os +from pathlib import Path +import time +import traceback + +from hf_space.gpu_mode import enabled + +_resident = None +_foreground_resident = None +_mounted_started = False + + +def start_mounted_app(): + """Run the SDK's launch hook for our mounted Gradio/FastAPI application. + + spaces 0.51.3 hooks Blocks.launch(), which a uvicorn-mounted app never calls. + All lease functions are registered before this hook; subsequent preload is + CPU-only, so no post-pack emulated CUDA tensors are introduced at startup. + """ + global _mounted_started + if enabled() and not _mounted_started: + from spaces.zero import startup + startup() + _mounted_started = True + + +def _gpu(setting, default): + seconds = int(os.environ.get(setting, str(default))) + if not 1 <= seconds <= 1200: + raise ValueError(setting + ' must be between 1 and 1200 seconds') + if enabled(): + import spaces + return spaces.GPU(duration=seconds) + return lambda function: function + + +def preload(config, notify, *, factory=None): + """Prepare host tensors before ZeroGPU forks its first GPU request.""" + global _resident + if factory is None: + from hf_space.resident_pipeline import ResidentPipeline + factory = ResidentPipeline + candidate = factory(config) + candidate.prepare(notify, cpu_only=True) + if not candidate.pool.sealed: + raise RuntimeError('The CPU model pool was not sealed') + _resident = candidate + return candidate.pool.summary() + + +def clear(): + global _resident + _resident = None + + +def preload_foreground(model_path, notify, *, factory=None): + """Load RMBG alone before the foreground wrapper forks its first worker.""" + global _foreground_resident + if factory is None: + from hf_space.resident_pipeline import ForegroundResident + factory = ForegroundResident + candidate = factory(model_path) + candidate.prepare(notify) + if not candidate.pool.sealed: + raise RuntimeError('The foreground CPU model pool was not sealed') + _foreground_resident = candidate + return candidate.pool.summary() + + +def _run(job_path, seed, foreground_prepared=False, *, foreground=False, stage='generate'): + resident = _foreground_resident if foreground else _resident + if resident is None: + raise RuntimeError('The CPU model pool is not ready') + from hf_space import runtime_worker + job = Path(job_path) + with (job / 'inference.log').open('w', encoding='utf-8', buffering=1) as log: + with redirect_stdout(log), redirect_stderr(log): + try: + if foreground: + runtime_worker.execute_foreground(resident, job) + elif stage == 'geometry': + runtime_worker.execute_geometry(resident, job, seed, + foreground_prepared=foreground_prepared) + elif stage == 'texture': + runtime_worker.execute_texture(resident, job, seed) + else: + runtime_worker.execute(resident, job, seed, + foreground_prepared=foreground_prepared) + except BaseException: + traceback.print_exc() + raise + finally: + resident.pool.restore_all() + return resident.pool.summary() + + +@_gpu('SUPRATEX_ZERO_GENERATE_SECONDS', 180) +def generate(job_path, seed, foreground_prepared=False): + return _run(job_path, seed, foreground_prepared) + + +@_gpu('SUPRATEX_ZERO_GEOMETRY_SECONDS', 180) +def geometry(job_path, seed, foreground_prepared=False): + return _run(job_path, seed, foreground_prepared, stage='geometry') + + +@_gpu('SUPRATEX_ZERO_TEXTURE_SECONDS', 180) +def texture(job_path, seed): + return _run(job_path, seed, stage='texture') + + +@_gpu('SUPRATEX_ZERO_FOREGROUND_SECONDS', 60) +def foreground(job_path): + return _run(job_path, 0, foreground=True) + + +@_gpu('SUPRATEX_ZERO_EXPORT_SECONDS', 180) +def export(source, output, reference, work, token): + from hf_space import extraction_progress, ovoxel_results + from hf_space.timing import observe_stages + events = [] + began = time.monotonic() + + def stage(name): + if name in extraction_progress.STAGES: + extraction_progress.stage(token, name) + events.append({'stage': name, 'start_seconds': round(time.monotonic() - began, 3)}) + + with observe_stages(lambda name: stage('export_gutter' if name == 'gutter' else name)): + ovoxel_results._export_local(Path(source), Path(output), gutter=reference is None) + if reference is not None: + from supratex.inference.reprojection import run_front_view + output = str(run_front_view(Path(output), Path(reference), Path(work), stage=stage, + keep=ovoxel_results._observation_keep(Path(reference)))) + return str(output), events + + +@_gpu('SUPRATEX_ZERO_RECOVER_SECONDS', 120) +def recover(identifier): + from hf_space.preview_recovery import _recover + _recover(identifier, render_html=False) + return identifier diff --git a/infer.py b/infer.py index 3fac13c5074cd0d15c19684cc6a4baab96fc24b7..ebeff01332a7d810f4bebeb1ca67a7e196f11f1c 100644 --- a/infer.py +++ b/infer.py @@ -116,6 +116,8 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--moge-model", help="local MoGe-2 checkpoint/directory or HF repo ID") parser.add_argument("--rmbg-model", help="local RMBG-2 directory or HF repo ID; required for RGB input") parser.add_argument("--use-input-alpha", action="store_true", help="require and reuse foreground RGBA alpha") + parser.add_argument("--no-filter", "--no-image-depth", dest="no_filter", action="store_true", + help="disable the default single-view image/geometry depth pre-filter; dual-view does not use it") parser.add_argument("--resolution-level", type=int, default=9, help="MoGe-2 resolution level") parser.add_argument("--geometry-package", type=Path, help="native or paper geometry source checkout; see docs/geometry.md") parser.add_argument("--geometry-model", help="local geometry weights directory or HF repo ID") @@ -289,13 +291,19 @@ def run(args: argparse.Namespace, *, resident=None, foreground_prepared: bool = from supratex.inference import run_completion work.mkdir(parents=True, exist_ok=True) handoff = work / "input" + filter_requested = not getattr(args, 'no_filter', False) + filter_enabled = filter_requested and not args.dual_view + filter_disabled_reason = ('user_disabled' if not filter_requested else + 'dual_view_unsupported' if args.dual_view else None) receipt = {"status": "running", "stage": "frontend", "input": str(image), "output": str(output), "work_dir": str(work), "seed": args.seed, "gutter": args.gutter, "material_backend": args.material_backend, "material_seed": 42, "runtime": runtime, "model_config": str(config), "checkpoint": str(checkpoint), "dual_view": args.dual_view, "refine": args.refine, "remove_floaters": args.remove_floaters, "preview_only": args.preview_only, - "foreground_prepared": foreground_prepared} + "foreground_prepared": foreground_prepared, + "image_depth_filter": {"requested": filter_requested, "enabled": filter_enabled, + "disabled_reason": filter_disabled_reason}} def record() -> None: temporary = work / ".run.json.tmp" temporary.write_text(json.dumps(receipt, indent=2) + "\n", encoding="utf-8") @@ -305,6 +313,7 @@ def run(args: argparse.Namespace, *, resident=None, foreground_prepared: bool = print(f"[1/2] Foreground, MoGe-2, geometry, material maps -> {handoff}", flush=True) result = run_frontend(image, handoff, moge_model=args.moge_model, rmbg_model=args.rmbg_model, use_input_alpha=args.use_input_alpha, resolution_level=args.resolution_level, + export_image_depth=filter_enabled, foreground_prepared=foreground_prepared, geometry_package=args.geometry_package.expanduser().resolve(), geometry_model=args.geometry_model, geometry_revision=args.geometry_revision, geometry_seed=args.seed, diff --git a/requirements.txt b/requirements.txt index b740f9cadfbea607eda9297de9377a3b9ee9182e..a56558fb9d98e3e7075f5248c501c53e856e5a24 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ torch==2.8.0 torchvision==0.23.0 gradio[oauth]==6.10.0 -spaces +spaces==0.51.3 # Recorded from the working Python 3.10 / PyTorch 2.6.0 CUDA 12.4 environment. # Install torch/torchvision and compiled extensions separately; see DEPENDENCIES.md. numpy==2.2.6 diff --git a/supratex/frontend/camera.py b/supratex/frontend/camera.py index 890c063b5d49dda47ea46820eac1d036fdc520bd..494b23db7bb5505328ef06f437bcc787d62ead2c 100644 --- a/supratex/frontend/camera.py +++ b/supratex/frontend/camera.py @@ -29,10 +29,37 @@ def camera_from_intrinsics(intrinsics: np.ndarray, width: int, height: int, *, r "moge_resolution_level": int(resolution_level)} +def _image_depth(result: dict, width: int, height: int) -> np.ndarray: + """Copy MoGe camera-z, rejecting outputs that changed the source pixel grid.""" + value = result.get("depth") + if value is None: + raise ValueError("MoGe did not return an image depth map") + depth = value.detach().float().cpu().numpy() + if depth.shape == (1, height, width): + depth = depth[0] + if depth.shape != (height, width): + raise ValueError("MoGe depth must preserve the input image pixel grid") + valid = np.isfinite(depth) & (depth > 0) + mask = result.get("mask") + if mask is not None: + mask = mask.detach().cpu().numpy() + if mask.shape == (1, height, width): + mask = mask[0] + if mask.shape != depth.shape: + raise ValueError("MoGe valid mask must match the image depth grid") + valid &= np.isfinite(mask) & (mask > 0) + return np.where(valid, depth, np.nan).astype(np.float32) + + def estimate(image: Path, model_path: Path, module_name: str = "moge.model.v2", resolution_level: int = 9, - device: str = "cuda", *, resident=None) -> dict: + device: str = "cuda", *, depth_output: Path | None = None, resident=None) -> dict: + """Estimate camera and optional image-aligned depth with the same model call.""" if not 0 <= resolution_level <= 9: raise ValueError("resolution_level must be in [0, 9]") + if depth_output is not None: + depth_output = Path(depth_output) + if depth_output.exists(): + raise FileExistsError(f"image depth output must be new: {depth_output}") module = importlib.import_module(module_name) model_cls = getattr(module, "MoGeModel", None) if model_cls is None: @@ -49,8 +76,14 @@ def estimate(image: Path, model_path: Path, module_name: str = "moge.model.v2", tensor = torch.from_numpy(np.asarray(rgb, dtype=np.float32) / 255.0).permute(2, 0, 1).to(device) with torch.inference_mode(): result = model.infer(tensor, resolution_level=resolution_level, force_projection=True, apply_mask=True) - return camera_from_intrinsics(result["intrinsics"].squeeze().float().cpu().numpy(), width, height, - resolution_level=resolution_level) + camera = camera_from_intrinsics(result["intrinsics"].squeeze().float().cpu().numpy(), width, height, + resolution_level=resolution_level) + if depth_output is not None: + depth = _image_depth(result, width, height) + depth_output.parent.mkdir(parents=True, exist_ok=True) + with depth_output.open("xb") as stream: + np.save(stream, depth, allow_pickle=False) + return camera finally: model.cpu() del tensor, result, model @@ -67,9 +100,11 @@ def main() -> None: help="MoGe-2 detail level from 0 (fastest) to 9 (highest).") parser.add_argument("--device", default="cuda") parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--depth-output", type=Path, help="save image-aligned camera-z as .npy; invalid pixels become NaN") args = parser.parse_args() args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(json.dumps(estimate(args.image, args.model, args.module, args.resolution_level, args.device), indent=2) + "\n") + args.output.write_text(json.dumps(estimate(args.image, args.model, args.module, args.resolution_level, args.device, + depth_output=args.depth_output), indent=2) + "\n") if __name__ == "__main__": diff --git a/supratex/frontend/pipeline.py b/supratex/frontend/pipeline.py index 5aa671e2fd633bb7c9d3001dd1b1f1e9f3c41f18..c6d2bc0baa3d4028fb9f28a9b3c52d4d7fff1100 100644 --- a/supratex/frontend/pipeline.py +++ b/supratex/frontend/pipeline.py @@ -112,12 +112,15 @@ def run(input_image: Path, output: Path, *, moge_model: str | Path, geometry_scale_search: bool = False, geometry_target_padding: int = 3, geometry_max_iterations: int = 2, geometry_seed: int = 42, record_material_provenance: bool = False, attention: str = "auto", low_vram: bool | None = None, - parallel_frontend: bool = False, resident=None, foreground_prepared: bool = False) -> dict[str, str]: + parallel_frontend: bool = False, resident=None, foreground_prepared: bool = False, + geometry_only: bool = False, export_image_depth: bool = True) -> dict[str, str]: """Build the condition, canonical camera, and material-map handoff in one directory.""" if resident is not None and (parallel_frontend or material_backend != "qwen-image"): raise ValueError("Resident frontend supports sequential Qwen generation only") if type(foreground_prepared) is not bool or (foreground_prepared and resident is None): raise ValueError('Prepared foreground is reserved for the trusted resident pipeline') + if geometry_only and (geometry_package is None or parallel_frontend): + raise ValueError('Geometry-only inference requires sequential geometry generation') if record_material_provenance and firered_manifest is not None: raise ValueError("material provenance does not support a custom FireRed manifest; " "omit it to bind the current condition and use matching prompts for both views") @@ -167,14 +170,20 @@ def run(input_image: Path, output: Path, *, moge_model: str | Path, rmbg_model = resolve_model_dir(rmbg_model, cache_dir=cache_dir, local_files_only=local_files_only) output.mkdir(parents=True, exist_ok=True) resident_options = {"resident": resident} if resident is not None else {} + from hf_space.timing import span estimator = RMBG2(rmbg_model, device=device, **resident_options) if rmbg_model and not foreground_prepared else None try: - rgba, rgb, foreground_source = prepare_image(input_image, output_dir=output, use_input_alpha=use_input_alpha, - background_model=estimator, foreground_prepared=foreground_prepared) + with span('frontend.preprocess'): + rgba, rgb, foreground_source = prepare_image(input_image, output_dir=output, use_input_alpha=use_input_alpha, + background_model=estimator, foreground_prepared=foreground_prepared) finally: if estimator is not None: estimator.close() - camera = estimate(rgb, moge_model, moge_module, resolution_level, device, **resident_options) + # Depth gating is single-view only; a dual-view front packet must omit it. + depth_path = output / "image_depth.npy" if export_image_depth and not record_material_provenance else None + with span('frontend.camera'): + camera = estimate(rgb, moge_model, moge_module, resolution_level, device, + depth_output=depth_path, **resident_options) camera["foreground_source"] = foreground_source camera_path = output / "camera.json" camera_path.write_text(json.dumps(camera, indent=2) + "\n", encoding="utf-8") @@ -196,7 +205,18 @@ def run(input_image: Path, output: Path, *, moge_model: str | Path, executor.shutdown(wait=False) print(f"Frontend: geometry on GPU {gpus[0]}, material on GPU {gpus[1]} in parallel", flush=True) else: - geometry_result = geometry_call() + with span('frontend.geometry'): + geometry_result = geometry_call() + if geometry_only: + geometry_path = output / 'geometry.pt' + geometry_path.write_bytes(Path(geometry_result['geometry']).read_bytes()) + camera_path.write_bytes(Path(geometry_result['camera']).read_bytes()) + (output / 'geometry_receipt.json').write_bytes(Path(geometry_result['receipt']).read_bytes()) + result = {'condition_rgba': str(rgba), 'condition_rgb': str(rgb), 'camera': str(camera_path), + 'geometry': str(geometry_path), 'foreground_source': foreground_source} + if depth_path is not None: + result['image_depth'] = str(depth_path) + return result material_dir = output / "material" try: albedo, mra = predict_material_maps( @@ -228,6 +248,8 @@ def run(input_image: Path, output: Path, *, moge_model: str | Path, result = {"condition_rgba": str(rgba), "condition_rgb": str(rgb), "camera": str(camera_path), "albedo": str(final_albedo), "mra": str(final_mra), "foreground_source": foreground_source, "material_backend": material_backend} + if depth_path is not None: + result["image_depth"] = str(depth_path) if geometry_result: result.update(geometry=str(geometry_path), geometry_receipt=geometry_result["receipt"]) return result @@ -267,6 +289,8 @@ def main() -> None: parser.add_argument("--material-checkpoint-manifest", type=Path) parser.add_argument("--key", default="input") parser.add_argument("--record-material-provenance", action="store_true", help="write material bindings for optional dual-view inference") + parser.add_argument("--no-filter", "--no-image-depth", dest="no_image_depth", action="store_true", + help="disable the default depth pre-filter by omitting image_depth.npy (implied by --record-material-provenance)") parser.add_argument("--device", default="cuda") parser.add_argument("--attention", choices=("auto", "flash_attn_3", "flash_attn", "sdpa", "xformers"), default="auto") parser.add_argument("--low-vram", action=argparse.BooleanOptionalAction, default=None, @@ -290,6 +314,7 @@ def main() -> None: rmbg_model=args.rmbg_model, use_input_alpha=args.use_input_alpha, device=args.device, moge_module=args.moge_module, resolution_level=args.resolution_level, key=args.key, record_material_provenance=args.record_material_provenance, + export_image_depth=not args.no_image_depth, attention=args.attention, low_vram=args.low_vram), indent=2)) diff --git a/supratex/inference/image_depth.py b/supratex/inference/image_depth.py new file mode 100644 index 0000000000000000000000000000000000000000..181aa97116255c53bc9160e61bf8a3be48f48be8 --- /dev/null +++ b/supratex/inference/image_depth.py @@ -0,0 +1,115 @@ +"""Conservative image/geometry depth consistency for source-pixel admission. + +An auxiliary depth map supplies evidence about image-to-surface correspondence; +it is not ground truth. Invalid or degenerate fits leave the existing admission +rules unchanged and report why the additional gate was disabled. +""" +from __future__ import annotations + +import math + +import numpy as np + +IMAGE_DEPTH_GAP = 0.02 +IMAGE_DEPTH_HIGHPASS = 41 +IMAGE_DEPTH_BOUNDARY_RADIUS = 16 +MIN_FIT_PIXELS = 1000 + + +def image_depth_disagreement(depth: np.ndarray, image_depth: np.ndarray, + gap: float = IMAGE_DEPTH_GAP, + highpass: int = IMAGE_DEPTH_HIGHPASS, + boundary_radius: int = IMAGE_DEPTH_BOUNDARY_RADIUS, + stats: dict | None = None) -> np.ndarray: + """Return suspicious source pixels, using camera-z depths on the same grid. + + ``gap`` is an absolute canonical length (unit object cube). ``highpass`` + and ``boundary_radius`` are pixels of the observation grid, normally 1024. + The latter is a square/Chebyshev radius. Pixels without two finite positive + depths are never rejected. A source pixel is rejected only when its local + residual exceeds ``gap`` near a depth edge in the fitted image-depth map. + ``stats`` receives JSON-safe fit diagnostics, including disabled reasons. + """ + from scipy.ndimage import maximum_filter, median_filter, minimum_filter + + depth = np.asarray(depth) + image_depth = np.asarray(image_depth) + if depth.ndim != 2 or image_depth.shape != depth.shape: + raise ValueError("geometry and image depths must be matching 2D maps") + if not math.isfinite(gap) or gap <= 0: + raise ValueError("image-depth gap must be finite and positive") + if not isinstance(highpass, int) or highpass < 1 or highpass % 2 != 1: + raise ValueError("image-depth highpass must be a positive odd size") + if not isinstance(boundary_radius, int) or boundary_radius < 0: + raise ValueError("image-depth boundary radius must be nonnegative") + + valid_image = np.isfinite(image_depth) & (image_depth > 0) + both = np.isfinite(depth) & (depth > 0) & valid_image + reject = np.zeros(depth.shape, dtype=bool) + report = {"status": "disabled", "reason": "insufficient_valid_depth", + "valid_pixels": int(both.sum()), "gap": float(gap), + "highpass": highpass, "boundary_radius": boundary_radius, + "rejected_pixels": 0} + + def disabled(reason: str) -> np.ndarray: + report["reason"] = reason + if stats is not None: + stats.update(report) + return reject + + if report["valid_pixels"] < MIN_FIT_PIXELS: + return disabled("insufficient_valid_depth") + x, y = image_depth[both].astype(np.float64), depth[both].astype(np.float64) + keep = np.ones(x.size, dtype=bool) + for iteration in range(3): + count = int(keep.sum()) + if count < MIN_FIT_PIXELS: + return disabled("insufficient_fit_inliers") + xx, yy = x[keep], y[keep] + # Center the fit to avoid an ill-conditioned [large-depth, 1] matrix. + # A single depth layer cannot identify both scale and offset. + xmean, ymean = float(xx.mean()), float(yy.mean()) + span = float(np.ptp(xx)) + if span <= 1e-6 * max(float(np.max(np.abs(xx))), 1.0): + return disabled("degenerate_image_depth_span") + dx, dy = xx - xmean, yy - ymean + variance = float(np.mean(dx * dx)) + if variance <= np.finfo(np.float64).eps * max(xmean * xmean, 1.0): + return disabled("degenerate_image_depth_variance") + a = float(np.mean(dx * dy) / variance) + b = ymean - a * xmean + if not math.isfinite(a) or not math.isfinite(b) or a <= 0: + return disabled("invalid_affine_fit") + residual = np.abs(y - (a * x + b)) + if not np.isfinite(residual).all(): + return disabled("nonfinite_affine_residual") + report.update(fit_iterations=iteration + 1, fit_pixels=count, + scale=a, offset=b, + median_absolute_residual=float(np.median(residual))) + if iteration < 2: + # Include percentile ties. Strict '<' can discard every pixel of + # an exactly aligned plane, or all of one of two depth layers. + threshold = float(np.percentile(residual, 80)) + tolerance = 64 * np.finfo(np.float64).eps * max(float(np.max(np.abs(y))), 1.0) + keep = residual <= threshold + tolerance + + fitted = np.full(depth.shape, np.nan, dtype=np.float32) + fitted[valid_image] = (a * image_depth[valid_image].astype(np.float64) + b).astype(np.float32) + if not np.isfinite(fitted[both]).all() or not (fitted[both] > 0).all(): + return disabled("invalid_fitted_depth") + full = np.zeros(depth.shape, dtype=np.float32) + full[both] = depth[both] - fitted[both] + smooth = median_filter(full, size=highpass) + maximum = maximum_filter(np.where(valid_image, fitted, -np.inf), size=3, + mode="constant", cval=-np.inf) + minimum = minimum_filter(np.where(valid_image, fitted, np.inf), size=3, + mode="constant", cval=np.inf) + with np.errstate(invalid="ignore"): + boundary = valid_image & np.isfinite(maximum) & np.isfinite(minimum) & ((maximum - minimum) > gap) + near = maximum_filter(boundary, size=2 * boundary_radius + 1, mode="constant", cval=0) + reject[both] = (np.abs((full - smooth)[both]) > gap) & near[both] + report.update(status="applied", reason=None, rejected_pixels=int(reject.sum()), + boundary_pixels=int(boundary.sum())) + if stats is not None: + stats.update(report) + return reject diff --git a/supratex/inference/native_observation.py b/supratex/inference/native_observation.py index 5b5a6e632fb4a2f9e523fb9fe486d70b0860b84b..4796e7d4ccaa1db015c408d2cc39ff7cf2c6294c 100644 --- a/supratex/inference/native_observation.py +++ b/supratex/inference/native_observation.py @@ -11,6 +11,8 @@ from typing import Any import numpy as np import torch +from .image_depth import image_depth_disagreement + RESOLUTION = 1024 LATENT_RESOLUTION = 64 @@ -233,10 +235,17 @@ def _admitted_pixels(depth: np.ndarray, missed: np.ndarray, focal: tuple[float, def pixel_first_occupied(camera: dict[str, Any], occupancy: np.ndarray, trusted: torch.Tensor, depth_edge_rtol: float = 0.04, - stats: dict | None = None) -> torch.Tensor: - """Map trusted camera pixels to their first occupied volume cell.""" + stats: dict | None = None, image_depth: np.ndarray | None = None, + depth_gate_out: dict | None = None) -> torch.Tensor: + """Map trusted pixels to first occupied cells, optionally excluding depth conflicts. + + With image depth, ``depth_gate_out['keep_mask']`` contains only the depth + gate's keep mask for material sampling and reprojection. Alpha and geometry + admission remain separate. Diagnostics are written to ``stats['image_depth']``. + """ if "extrinsic" in camera: - return pixel_first_occupied_cv(camera, occupancy, trusted, depth_edge_rtol, stats=stats) + return pixel_first_occupied_cv(camera, occupancy, trusted, depth_edge_rtol, stats=stats, + image_depth=image_depth, depth_gate_out=depth_gate_out) if trusted.shape != (RESOLUTION, RESOLUTION): raise ValueError("trusted alpha mask must be 1024 by 1024") if "camera_angle_x" not in camera or "distance" not in camera: @@ -261,6 +270,12 @@ def pixel_first_occupied(camera: dict[str, Any], occupancy: np.ndarray, trusted: admitted = _admitted_pixels(depth, missed, (focal, focal), depth_edge_rtol, stats=stats) retained = hit_mask.copy() retained[hit_mask] &= admitted[y[hit_mask], x[hit_mask]] + if image_depth is not None: + mismatch = _image_depth_mask(camera, occupancy, image_depth, stats, depth_gate_out) + rejected = retained & mismatch[y, x] + retained &= ~mismatch[y, x] + if stats is not None: + stats.update(rejected_image_depth=int(rejected.sum()), admitted=int(retained.sum())) hit = owners[retained] if not hit.size: raise RuntimeError("observation admission removed every first-hit observation") @@ -292,6 +307,53 @@ def camera_rays(camera: dict[str, Any], ys: np.ndarray, xs: np.ndarray) -> tuple return eye, directions, rotation +def first_hit_depth(camera: dict[str, Any], occupancy: np.ndarray) -> np.ndarray: + """Return full-grid first-hit camera-z for canonical and complete CV cameras. + + The canonical camera looks down world -z; its positive view depth is + ``eye_z - world_z``. CV cameras instead use the z component after their + world-to-camera rotation. Neither convention uses Euclidean ray length. + """ + ys, xs = np.mgrid[0:RESOLUTION, 0:RESOLUTION] + ys, xs = ys.ravel(), xs.ravel() + if "extrinsic" in camera: + eye, directions, rotation = camera_rays(camera, ys, xs) + else: + angle = float(camera["camera_angle_x"]) + distance = float(camera["distance"]) * float(camera.get("mesh_scale", 1.0)) + if not math.isfinite(angle) or not 0 < angle < math.pi or not math.isfinite(distance) or distance <= 0: + raise ValueError("canonical camera requires a finite FOV in (0, pi) and positive effective distance") + # Keep the arithmetic order of the canonical admission ray builder. + focal = 16.0 / math.tan(angle / 2.0) * RESOLUTION / 32.0 + eye = np.array([0., 0., distance], dtype=np.float64) + directions = np.column_stack(((xs + .5 - RESOLUTION / 2) / focal, + -(ys + .5 - RESOLUTION / 2) / focal, + -np.ones(xs.size))).astype(np.float64) + rotation = np.diag([1., -1., -1.]) + owners, reasons, _ = _dda_batch_kernel()(eye, directions, occupancy, -.5, .5, RESOLUTION) + hits = reasons == 0 + centers = (owners[hits].astype(np.float64) + .5) / RESOLUTION - .5 + values = ((centers - eye) @ rotation.T)[:, 2] + if not np.isfinite(values).all() or not (values > 0).all(): + raise RuntimeError("first-hit owners must lie in front of the camera") + result = np.full((RESOLUTION, RESOLUTION), np.nan, dtype=np.float32) + result[ys[hits], xs[hits]] = values.astype(np.float32) + return result + + +def _image_depth_mask(camera: dict[str, Any], occupancy: np.ndarray, image_depth: np.ndarray, + stats: dict | None, depth_gate_out: dict | None) -> np.ndarray: + if np.asarray(image_depth).shape != (RESOLUTION, RESOLUTION): + raise ValueError("image depth must be an observation-resolution map") + diagnostic: dict[str, Any] = {} + mismatch = image_depth_disagreement(first_hit_depth(camera, occupancy), image_depth, stats=diagnostic) + if stats is not None: + stats["image_depth"] = diagnostic + if depth_gate_out is not None: + depth_gate_out["keep_mask"] = ~mismatch + return mismatch + + def pixel_hits_cv(camera: dict[str, Any], occupancy: np.ndarray, mask: np.ndarray) -> np.ndarray: """Return the raw first-hit silhouette, without trusted-mask depth filtering.""" if mask.shape != (RESOLUTION, RESOLUTION) or mask.dtype != np.bool_: @@ -305,7 +367,8 @@ def pixel_hits_cv(camera: dict[str, Any], occupancy: np.ndarray, mask: np.ndarra def pixel_first_occupied_cv(camera: dict[str, Any], occupancy: np.ndarray, trusted: torch.Tensor, - depth_edge_rtol: float = .04, stats: dict | None = None) -> torch.Tensor: + depth_edge_rtol: float = .04, stats: dict | None = None, + image_depth: np.ndarray | None = None, depth_gate_out: dict | None = None) -> torch.Tensor: """Apply the production first-hit admission rules to a full CV camera.""" if trusted.shape != (RESOLUTION, RESOLUTION) or trusted.dtype != torch.bool: raise ValueError("trusted alpha must be a boolean observation-resolution image") @@ -330,6 +393,12 @@ def pixel_first_occupied_cv(camera: dict[str, Any], occupancy: np.ndarray, trust intrinsic = np.asarray(camera.get("projection_intrinsic", camera.get("intrinsic")), dtype=np.float64) focal = (float(intrinsic[0, 0]) * RESOLUTION, float(intrinsic[1, 1]) * RESOLUTION) retained = hits & _admitted_pixels(depth, missed, focal, depth_edge_rtol, stats=stats)[ys, xs] + if image_depth is not None: + mismatch = _image_depth_mask(camera, occupancy, image_depth, stats, depth_gate_out) + rejected = retained & mismatch[ys, xs] + retained &= ~mismatch[ys, xs] + if stats is not None: + stats.update(rejected_image_depth=int(rejected.sum()), admitted=int(retained.sum())) kept = owners[retained].astype(np.int64, copy=False) if not kept.size: raise RuntimeError("observation admission removed every first-hit observation") diff --git a/supratex/inference/pipeline.py b/supratex/inference/pipeline.py index 1d2b7678416a367792b4a04dda51d66f4bd5ede4..91bab479944dab94005e474a43426852814bffad 100644 --- a/supratex/inference/pipeline.py +++ b/supratex/inference/pipeline.py @@ -196,6 +196,17 @@ def _load_case(input_dir: str | Path, torch: Any, image: Any, remove_floaters: b raise ValueError("camera.json must be valid JSON") from error if "camera_angle_x" not in camera or "distance" not in camera: raise ValueError("camera.json must provide canonical camera_angle_x and distance") + image_depth = None + if (root / "image_depth.npy").is_file(): + import numpy as np + raw = np.load(root / "image_depth.npy", allow_pickle=False) + if raw.shape != rgba.size[::-1]: + raise ValueError("image_depth.npy must match the condition_rgba.png pixel grid") + if raw.dtype.kind not in "fiu": + raise ValueError("image_depth.npy must contain numeric camera-z depths") + raw = raw.astype(np.float32) + raw[~np.isfinite(raw) | (raw <= 0)] = np.nan + image_depth = np.asarray(image.fromarray(raw).resize((RESOLUTION, RESOLUTION), image.Resampling.NEAREST), dtype=np.float32) if remove_floaters: import numpy as np from .geometry_cleanup import remove_floaters as clean @@ -203,7 +214,8 @@ def _load_case(input_dir: str | Path, torch: Any, image: Any, remove_floaters: b vertices, faces = vertices.contiguous(), faces.contiguous() print(f"Floaters: removed {report['removed_components']} of {report['components']} components " f"({report['removed_faces']} of {report['faces']} faces)", flush=True) - return {"vertices": vertices, "faces": faces, "rgba": rgba, "albedo": albedo, "mra": mra, "camera": camera} + return {"vertices": vertices, "faces": faces, "rgba": rgba, "albedo": albedo, "mra": mra, "camera": camera, + "image_depth": image_depth} def _runtime() -> dict[str, Any]: @@ -266,7 +278,8 @@ def _alignment_report(uv: Any, alpha: Any, admission: Mapping[str, int], np: Any rejected = [f"missed geometry {admission['trusted_rays'] - admission['first_hits']}", f"depth edge {admission['rejected_depth_edge']}"] rejected += [f"{label} {admission[key]}" for key, label in (("rejected_silhouette", "silhouette margin"), - ("rejected_grazing", "grazing")) if admission[key]] + ("rejected_grazing", "grazing"), + ("rejected_image_depth", "image-depth mismatch")) if admission.get(key)] print(f"Observation: silhouette IoU {report['silhouette_iou']:.3f} " f"(geometry inside image {report['geometry_inside_image']:.3f}, image covered {report['image_covered_by_geometry']:.3f}); " f"admitted {admission['admitted']} of {admission['trusted_rays']} trusted pixels ({', '.join(rejected)})", flush=True) @@ -276,7 +289,21 @@ def _alignment_report(uv: Any, alpha: Any, admission: Mapping[str, int], np: Any return report -def _partial(surface: Any, shape: Any, case: Mapping[str, Any], external: Any, projection: Any, encoder: Any, normalization: Mapping[str, Any], runtime: Mapping[str, Any], device: str) -> Any: +def _source_footprint_safe(keep: Any, grid: Any, runtime: Mapping[str, Any], device: str) -> Any: + """Reject material samples with nonzero weight in a depth-mismatch pixel. + + Use the actual material-sampling grid and interpolation conventions, rather + than a separately rounded CPU projection of voxel centres. + """ + torch, F, np = runtime["torch"], runtime["F"], runtime["np"] + unsafe = torch.from_numpy(np.asarray(~keep, dtype=np.float32)).to(device)[None, None] + contribution = F.grid_sample(unsafe, grid, mode="bilinear", align_corners=False, + padding_mode="border")[0, 0, 0] + return contribution <= 1e-6 + + +def _partial(surface: Any, shape: Any, case: Mapping[str, Any], external: Any, projection: Any, encoder: Any, normalization: Mapping[str, Any], runtime: Mapping[str, Any], device: str) -> tuple[Any, dict[str, Any]]: + """Encode observations and return the depth-gate handoff for front refine.""" torch, F, np, image, image_filter, sparse = runtime["torch"], runtime["F"], runtime["np"], runtime["Image"], runtime["ImageFilter"], runtime["sparse"] coords = surface.coords[:, 1:].long(); points = (coords.float() + .5) / RESOLUTION - .5 camera = (torch.cat((points, torch.ones_like(points[:, :1])), dim=1) @ external.T)[:, :3] @@ -289,13 +316,35 @@ def _partial(surface: Any, shape: Any, case: Mapping[str, Any], external: Any, p mask = image.fromarray(np.where(alpha >= 254, 255, 0).astype(np.uint8)).resize((RESOLUTION, RESOLUTION), image.Resampling.NEAREST).filter(image_filter.MinFilter(5)) trusted = torch.from_numpy(np.asarray(mask, dtype=np.uint8).copy() == 255) bits, ordered, order = native_observation.build_occupancy_index(coords) - admission: dict[str, int] = {} - owners = native_observation.pixel_first_occupied(dict(case["camera"]), bits, trusted, stats=admission) + admission: dict[str, Any] = {} + depth_gate: dict[str, Any] = {} + completion_metadata: dict[str, Any] = {} + if case.get("image_depth") is None: + owners = native_observation.pixel_first_occupied(dict(case["camera"]), bits, trusted, stats=admission) + else: + owners = native_observation.pixel_first_occupied(dict(case["camera"]), bits, trusted, stats=admission, + image_depth=case["image_depth"], depth_gate_out=depth_gate) _alignment_report(uv, alpha, admission, np, image, image_filter) + if admission.get("image_depth"): + diagnostic = admission["image_depth"] + detail = f" ({diagnostic['reason']})" if diagnostic.get("reason") else "" + print(f"Image depth gate: {diagnostic['status']}{detail}; " + f"{diagnostic['valid_pixels']} paired depth pixels, " + f"{diagnostic['rejected_pixels']} mismatch pixels", flush=True) visible = native_observation.owner_rows(coords, owners, ordered, order) + if "keep_mask" in depth_gate: + depth_keep = np.asarray(depth_gate["keep_mask"], dtype=bool) + safe = _source_footprint_safe(depth_keep, grid, runtime, device) + admission["rejected_image_depth_footprint_owners"] = int((visible & ~safe).sum()) + visible &= safe + admission["observed_voxels"] = int(visible.sum()) + completion_metadata = {"observation_keep": depth_keep, "observation_admission": admission} + print(f"Observation image depth: excluded {admission['rejected_image_depth_footprint_owners']} " + f"additional voxel samples crossing the mismatch mask; {admission['observed_voxels']} observed voxels", flush=True) if not bool(visible.any()): raise RuntimeError("the trusted alpha gate leaves no observed material surface") with torch.no_grad(): - return native_observation.encode_partial(encoder, coords, visible, sampled(case["albedo"]), sampled(case["mra"]), shape.coords[:, 1:], dict(normalization), sparse, device) + observation = native_observation.encode_partial(encoder, coords, visible, sampled(case["albedo"]), sampled(case["mra"]), shape.coords[:, 1:], dict(normalization), sparse, device) + return observation, completion_metadata def _export(output: Path, representation: Any, texture_size: int, decimation_target: int, texture_format: str, runtime: Mapping[str, Any], *, gutter: bool = False, low_vram: bool = False) -> None: @@ -372,6 +421,31 @@ def _completion_stage(model: Any, torch: Any, device: str, *, layer_offload: boo torch.cuda.empty_cache() +def _save_observation_metadata(output: str | Path, metadata: Mapping[str, Any]) -> None: + """Keep the actual completion gate for later preview recovery and export. + + The sidecars belong to this result, never to the immutable geometry input. + Legacy inputs without image depth deliberately produce no sidecars. + """ + keep = metadata.get("observation_keep") + if keep is None: + return + import numpy as np + keep = np.asarray(keep) + if keep.shape != (RESOLUTION, RESOLUTION) or keep.dtype != np.bool_: + raise ValueError("observation keep must be a boolean observation-resolution mask") + admission = metadata.get("observation_admission") + if not isinstance(admission, dict): + raise ValueError("observation admission must contain the depth-gate diagnostics") + payload = json.dumps(admission, indent=2, allow_nan=False) + "\n" + output = Path(output) + output.parent.mkdir(parents=True, exist_ok=True) + with output.with_suffix(".observation_keep.npy").open("xb") as stream: + np.save(stream, keep, allow_pickle=False) + with output.with_suffix(".observation_admission.json").open("x", encoding="utf-8") as stream: + stream.write(payload) + + def _run_completion_low_vram(*, case: Mapping[str, Any], back_case: Mapping[str, Any] | None, ckpts: Path, config: Mapping[str, Any], dataset: Mapping[str, Any], checkpoint_path: Path, partial_stats: Mapping[str, Any], resolved: Any, @@ -390,15 +464,15 @@ def _run_completion_low_vram(*, case: Mapping[str, Any], back_case: Mapping[str, del shape_encoder texture_encoder = (resident.completion_model("texture_encoder") if resident is not None else _load_vae(models, _stem(ckpts / "tex_enc_next_dc_f16c32_fp16", "material encoder"))) - dual_stats = {} + completion_metadata = {} with _completion_stage(texture_encoder, torch, device): if back_case is None: external, projection = _camera(case["camera"], device, torch, runtime["utils3d"]) - observation = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) + observation, completion_metadata = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) del external, projection else: from .dual_view import encode_two_views - observation, dual_stats = encode_two_views(surface, shape, case, back_case, texture_encoder, + observation, completion_metadata = encode_two_views(surface, shape, case, back_case, texture_encoder, partial_stats, runtime, device) del texture_encoder, surface torch.cuda.empty_cache() @@ -448,12 +522,13 @@ def _run_completion_low_vram(*, case: Mapping[str, Any], back_case: Mapping[str, torch.cuda.empty_cache() representation = runtime["MeshWithVoxel"](case["vertices"], case["faces"], [-.5, -.5, -.5], 1. / RESOLUTION, decoded_cpu.coords[:, 1:].int(), decoded_cpu.feats.float().clamp(0, 1), torch.Size([1, 6, RESOLUTION, RESOLUTION, RESOLUTION]), layout=LAYOUT) del decoded_cpu + _save_observation_metadata(output, completion_metadata) if preview_only: from hf_space.ovoxel_preview import save_and_render save_and_render(Path(output), representation.to(device)) else: _export(Path(output), representation, texture_size, decimation_target, texture_format, runtime, gutter=gutter, low_vram=True) - return dual_stats + return completion_metadata def _run_completion(*, input_dir: str | Path, output: str | Path, model_config: str | Path, checkpoint: str | Path, trellis_root: str | Path, seed: int = 42, device: str = "cuda", texture_size: int = 4096, decimation_target: int = 500000, texture_format: str = "png", gutter: bool = False, back_case: Mapping[str, Any] | None = None, remove_floaters: bool = False, preview_only: bool = False, resident=None) -> dict[str, Any]: @@ -470,6 +545,9 @@ def _run_completion(*, input_dir: str | Path, output: str | Path, model_config: runtime = _runtime(); torch, models = runtime["torch"], runtime["models"] if not torch.cuda.is_available(): raise RuntimeError("CUDA is unavailable") case = _load_case(input_dir, torch, runtime["Image"], remove_floaters); ckpts = _completion_vae_root(trellis_root) + if back_case is not None and (case.get("image_depth") is not None or back_case.get("image_depth") is not None): + raise ValueError("image_depth.npy consistency gating currently supports single-view completion only; " + "remove it or prepare the front view with --no-image-depth") if os.environ.get("SUPRATEX_LOW_VRAM", "0") == "1": return _run_completion_low_vram(case=case, back_case=back_case, ckpts=ckpts, config=config, dataset=dataset, checkpoint_path=checkpoint_path, partial_stats=partial_stats, @@ -482,12 +560,12 @@ def _run_completion(*, input_dir: str | Path, output: str | Path, model_config: external, projection = _camera(case["camera"], device, torch, runtime["utils3d"]) shape, surface = _shape(case["vertices"], case["faces"], shape_encoder, runtime, device) shape_encoder.to("cpu"); torch.cuda.empty_cache() - dual_stats = {} + completion_metadata = {} if back_case is None: - observation = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) + observation, completion_metadata = _partial(surface, shape, case, external, projection, texture_encoder, partial_stats, runtime, device) else: from .dual_view import encode_two_views - observation, dual_stats = encode_two_views(surface, shape, case, back_case, texture_encoder, + observation, completion_metadata = encode_two_views(surface, shape, case, back_case, texture_encoder, partial_stats, runtime, device) texture_encoder.to("cpu"); torch.cuda.empty_cache() from supratex.pipelines.observation_guided_completion import build_completion_model, sample_material_latents @@ -513,12 +591,13 @@ def _run_completion(*, input_dir: str | Path, output: str | Path, model_config: with torch.no_grad(): _, guide = shape_decoder(shape, return_subs=True); decoded = material_decoder(generated, guide_subs=guide) * .5 + .5 if decoded.feats.shape[1] != 6: raise RuntimeError("material decoder must emit six PBR channels") representation = runtime["MeshWithVoxel"](case["vertices"].to(device), case["faces"].to(device), [-.5, -.5, -.5], 1. / RESOLUTION, decoded.coords[:, 1:].int(), decoded.feats.float().clamp(0, 1), torch.Size([1, 6, RESOLUTION, RESOLUTION, RESOLUTION]), layout=LAYOUT) + _save_observation_metadata(output, completion_metadata) if preview_only: from hf_space.ovoxel_preview import save_and_render save_and_render(Path(output), representation.to(device)) else: _export(Path(output), representation, texture_size, decimation_target, texture_format, runtime, gutter=gutter) - return dual_stats + return completion_metadata def _require_reprojection_atlas(texture_size: int, texture_format: str) -> None: @@ -585,7 +664,7 @@ def _run_front_reprojection(arguments: dict[str, Any], work_dir: str | Path | No write_json(work / "receipt.json", receipt) try: native = work / "preds/mesh.glb" - _run_completion(**{**arguments, "input_dir": source, "output": native, "gutter": False}) + completion = _run_completion(**{**arguments, "input_dir": source, "output": native, "gutter": False}) # Reprojection runs in separate CUDA workers; release cached blocks first. import gc import torch @@ -595,7 +674,13 @@ def _run_front_reprojection(arguments: dict[str, Any], work_dir: str | Path | No write_json(work / "receipt.json", receipt) from .reprojection import run_front_view try: - final = run_front_view(native, source, work / "reproject") + keep = (completion or {}).get("observation_keep") + if keep is None: + final = run_front_view(native, source, work / "reproject") + else: + receipt["image_depth_admission"] = completion.get("observation_admission", {}) + receipt["reprojection_excluded_source_pixels"] = int((~keep).sum()) + final = run_front_view(native, source, work / "reproject", keep=keep) except Exception as error: _publish_unrefined(native, output, arguments["gutter"], error) receipt.update(status="valid", stage="complete", refine="skipped", @@ -620,7 +705,7 @@ def run_completion(*, input_dir: str | Path, output: str | Path, model_config: s decimation_target: int = 500000, texture_format: str = "png", gutter: bool = False, back_view_dir: str | Path | None = None, work_dir: str | Path | None = None, refine: bool = False, remove_floaters: bool = False, preview_only: bool = False, - resident=None) -> None: + resident=None) -> dict[str, Any] | None: """Complete one view by default, optionally refining it, or publish a refined dual-view result.""" arguments = dict(input_dir=input_dir, output=output, model_config=model_config, checkpoint=checkpoint, trellis_root=trellis_root, seed=seed, device=device, @@ -632,6 +717,10 @@ def run_completion(*, input_dir: str | Path, output: str | Path, model_config: s arguments["resident"] = resident if preview_only and (back_view_dir is not None or refine): raise ValueError("o-voxel preview must precede single-view mesh export") + if back_view_dir is not None and any((Path(root) / "image_depth.npy").is_file() + for root in (input_dir, back_view_dir)): + raise ValueError("image_depth.npy consistency gating currently supports single-view completion only; " + "remove it or prepare the front view with --no-image-depth") # Dual-view refines both views itself, so --refine needs no separate pass there. if refine and back_view_dir is None: _run_front_reprojection(arguments, work_dir) @@ -639,8 +728,7 @@ def run_completion(*, input_dir: str | Path, output: str | Path, model_config: s if back_view_dir is None: if work_dir is not None: raise ValueError("completion --work-dir is reserved for dual-view inference or --refine") - _run_completion(**arguments) - return + return _run_completion(**arguments) from .dual_view import CAMERA_PROTOCOL, PROTOCOL, digest, load_back_case, write_json output = Path(output).expanduser().resolve() diff --git a/supratex/inference/reprojection/__init__.py b/supratex/inference/reprojection/__init__.py index ca778aba7655315abadb73f757b3641046de54fb..067c7f6772e8a9f2c4c9402114cbcf41dc25c9aa 100644 --- a/supratex/inference/reprojection/__init__.py +++ b/supratex/inference/reprojection/__init__.py @@ -32,7 +32,7 @@ PROFILE = "supratex.dual_view.reproject_continuous" WORKER_TIMEOUT_SECONDS = 1800 -def _prepare_inputs(source: Path, destination: Path, camera: dict[str, Any]) -> dict[str, Any]: +def _prepare_inputs(source: Path, destination: Path, camera: dict[str, Any], keep: Any = None) -> dict[str, Any]: """Align observation alpha and RGB once, preserving native albedo/MRA maps.""" from PIL import Image @@ -47,6 +47,11 @@ def _prepare_inputs(source: Path, destination: Path, camera: dict[str, Any]) -> resized = rgb.convert("RGB").resize((1024, 1024), Image.Resampling.LANCZOS) alpha = rgba.convert("RGBA").getchannel("A").point(lambda value: 255 if value >= 254 else 0) alpha = alpha.resize((1024, 1024), Image.Resampling.NEAREST) + if keep is not None: + import numpy as np + # Apply at the actual projection grid. A 1024 -> small source -> + # 1024 round trip could otherwise resurrect thin rejected strips. + alpha = Image.fromarray(np.where(keep, np.asarray(alpha), 0).astype(np.uint8)) resized.save(destination / "condition_rgb.png") aligned = resized.copy() aligned.putalpha(alpha) @@ -61,6 +66,27 @@ def _prepare_inputs(source: Path, destination: Path, camera: dict[str, Any]) -> def _command(module: str, arguments: list[str], log: Path) -> None: + from hf_space.gpu_mode import enabled as zero_gpu_enabled + if zero_gpu_enabled(): + # ZeroGPU's allocated process owns the CUDA lease. A new interpreter + # would lose it; call each existing stage in that process instead. + from contextlib import redirect_stderr, redirect_stdout + import gc + import importlib + if module not in ('_project', '_continuous'): + raise ValueError('Unknown reprojection stage') + previous = sys.argv + try: + sys.argv = [module, *arguments] + with log.open('x', encoding='utf-8') as stream: + with redirect_stdout(stream), redirect_stderr(stream): + importlib.import_module(f'supratex.inference.reprojection.{module}').main() + finally: + sys.argv = previous + gc.collect() + import torch + torch.cuda.empty_cache() + return command = [sys.executable, "-m", f"supratex.inference.reprojection.{module}", *arguments] # The web process adds bundled runtimes such as utils3d to sys.path after # startup. Python subprocesses do not inherit those additions automatically. @@ -87,13 +113,14 @@ def _command(module: str, arguments: list[str], log: Path) -> None: raise RuntimeError(f"reprojection worker failed with exit code {result}; see {log}") -def _run_side(baseline: Path, source: Path, output: Path, camera: dict[str, Any], *, stage=None) -> Path: +def _run_side(baseline: Path, source: Path, output: Path, camera: dict[str, Any], keep: Any = None, *, stage=None) -> Path: """Publish one view only after the projection and continuous candidates pass.""" output.mkdir(parents=True) baseline_binding = bind_file(baseline) code = {path.name: bind_file(path) for path in Path(__file__).parent.glob("*.py")} inputs = output / "inputs_1024" - source_bindings = _prepare_inputs(source, inputs, camera) + source_bindings = (_prepare_inputs(source, inputs, camera) if keep is None else + _prepare_inputs(source, inputs, camera, keep=keep)) input_bindings = {path.name: bind_file(path) for path in inputs.iterdir() if path.is_file()} renderer = Path(__file__).with_name("_renderer.py").resolve() project, continuous = output / "project", output / "continuous" @@ -201,13 +228,17 @@ def _validate_side(output: Path) -> dict[str, Any]: return receipt -def run_front_view(native_glb: Path, input_dir: Path, output_dir: Path, *, stage=None) -> Path: +def run_front_view(native_glb: Path, input_dir: Path, output_dir: Path, keep: Any = None, *, stage=None) -> Path: """Reproject the single reference view onto a single-view result. This is the front half of dual-view texturing with the canonical reference camera from ``camera.json``. Inputs are staged under ``output_dir`` so their bindings stay valid; a missing ``condition_rgb.png`` is derived from the RGBA observation with black premultiplication, as the frontend writes it. + Optional ``keep`` excludes source pixels identified as image/geometry depth + mismatches. It changes only staged projection trust, not source RGB, source + files, or the final mesh's opacity. Other observation-admission filters are + not passed through this mask. """ import numpy as np from PIL import Image @@ -217,6 +248,10 @@ def run_front_view(native_glb: Path, input_dir: Path, output_dir: Path, *, stage if output.exists(): raise FileExistsError(f"front reprojection requires a new destination: {output}") camera = read_json(source / "camera.json") + if keep is not None: + keep = np.asarray(keep) + if keep.shape != (1024, 1024) or keep.dtype != np.bool_: + raise ValueError("reprojection keep must be a boolean 1024 by 1024 source-image mask") try: camera = {"camera_angle_x": float(camera["camera_angle_x"]), "distance": float(camera["distance"]), "mesh_scale": float(camera.get("mesh_scale", 1.0))} @@ -232,6 +267,8 @@ def run_front_view(native_glb: Path, input_dir: Path, output_dir: Path, *, stage with Image.open(source / "condition_rgba.png") as image: rgba = np.asarray(image.convert("RGBA"), dtype=np.float32) / 255. Image.fromarray((np.clip(rgba[..., :3] * rgba[..., 3:4], 0, 1) * 255).astype(np.uint8)).save(staged / "condition_rgb.png") + if keep is not None: + return _run_side(native, staged, output / "front", camera, keep=keep, stage=stage) return _run_side(native, staged, output / "front", camera, stage=stage) diff --git a/tests/test_depth_prefilter.py b/tests/test_depth_prefilter.py new file mode 100644 index 0000000000000000000000000000000000000000..381476a15d6bd0116b60a5651bddb8502e4ade35 --- /dev/null +++ b/tests/test_depth_prefilter.py @@ -0,0 +1,379 @@ +"""Depth filtering changes material observations, never the saved geometry.""" +import importlib.util +import json +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import numpy as np +from PIL import Image, ImageFilter +import pytest + +torch = pytest.importorskip('torch') +pytest.importorskip('scipy') + + +@pytest.fixture +def core(): + # Avoid the root package's unrelated eager CUDA-extension imports while + # exercising the real NumPy/SciPy/Torch inference functions on the CPU. + name = '_depth_prefilter_test' + root = Path(__file__).resolve().parents[1]/'supratex'/'inference' + package = ModuleType(name) + package.__path__ = [str(root)] + sys.modules[name] = package + spec = importlib.util.spec_from_file_location(name+'.pipeline', root/'pipeline.py') + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + try: + yield module + finally: + for key in list(sys.modules): + if key == name or key.startswith(name+'.'): + sys.modules.pop(key, None) + + +def depth_pair(): + _, x = np.mgrid[:128, :128] + image = (1 + .001*x + .3*(x >= 64)).astype(np.float32) + geometry = (2*image + .4).astype(np.float32) + return geometry, image + + +def test_affine_depth_alignment_keeps_observations_and_inputs_exact(core): + from _depth_prefilter_test.image_depth import image_depth_disagreement + geometry, image = depth_pair() + originals = geometry.tobytes(), image.tobytes() + stats = {} + rejected = image_depth_disagreement(geometry, image, stats=stats) + assert rejected.dtype == np.bool_ and not rejected.any() + assert stats['status'] == 'applied' and stats['scale'] == pytest.approx(2, abs=1e-5) + assert stats['offset'] == pytest.approx(.4, abs=1e-5) + assert originals == (geometry.tobytes(), image.tobytes()) + + +def test_shifted_occlusion_boundary_rejects_only_nearby_depth_mismatches(core): + from _depth_prefilter_test.image_depth import image_depth_disagreement + geometry, image = depth_pair() + geometry[:, 64:72] -= .6 + geometry[30:36, 10:16] += .5 # Depth disagreement far from an image edge. + image[10, 65] = np.nan + stats = {} + rejected = image_depth_disagreement(geometry, image, stats=stats) + assert rejected[:, 66:71].mean() > .95 + assert not rejected[:, :40].any() and not rejected[:, 90:].any() + assert not rejected[10, 65] + assert stats['status'] == 'applied' and stats['rejected_pixels'] == int(rejected.sum()) + + +@pytest.mark.parametrize('reason', ['insufficient_valid_depth', 'degenerate_image_depth_span']) +def test_unreliable_depth_fit_disables_only_the_new_gate(core, reason): + from _depth_prefilter_test.image_depth import image_depth_disagreement + geometry, image = np.ones((64, 64), dtype=np.float32), np.ones((64, 64), dtype=np.float32) + if reason == 'insufficient_valid_depth': + image[:60] = np.nan + stats = {} + assert not image_depth_disagreement(geometry, image, stats=stats).any() + assert stats['status'] == 'disabled' and stats['reason'] == reason + + +def test_bilinear_footprint_cannot_blend_a_rejected_neighbor(core): + keep = np.ones((4, 4), dtype=bool) + keep[1, 1] = False + # Pixel centers (1,1), (2,1), and the half-way point between them. + grid = torch.tensor([[[[-.25, -.25], [.25, -.25], [0., -.25]]]]) + safe = core._source_footprint_safe(keep, grid, {'torch':torch, 'F':torch.nn.functional, 'np':np}, 'cpu') + assert safe.tolist() == [False, True, False] + assert not keep[1, 1] and keep.sum() == 15 + + +@pytest.mark.parametrize('cv_camera', [False, True]) +def test_full_grid_geometry_depth_uses_camera_z_for_both_camera_conventions(core, monkeypatch, cv_camera): + native = core.native_observation + monkeypatch.setattr(native, 'RESOLUTION', 4) + calls = [] + def kernel(eye, directions, occupancy, low, high, resolution): + calls.append((eye.copy(), directions.copy())) + assert len(directions) == 16 and resolution == 4 + return np.tile([1, 1, 1], (16, 1)), np.zeros(16, np.uint8), np.zeros(16, np.int64) + monkeypatch.setattr(native, '_dda_batch_kernel', lambda: kernel) + camera = {'camera_angle_x':.7, 'distance':3.} + if cv_camera: + camera = {'extrinsic':[[1,0,0,0],[0,-1,0,0],[0,0,-1,3],[0,0,0,1]], + 'intrinsic':[[1,0,.5],[0,1,.5],[0,0,1]]} + depth = native.first_hit_depth(camera, np.zeros(1, np.uint8)) + np.testing.assert_array_equal(depth, np.full((4,4), 3.125, np.float32)) + assert len(calls) == 1 + + +def case_files(folder, size=8): + folder.mkdir() + vertices = torch.tensor([[0., -.0, .25], [.5, -.5, .5], [.125, .25, -.375]]) + faces = torch.tensor([[2, 0, 1]], dtype=torch.int64) + torch.save({'verts':vertices, 'faces':faces}, folder/'geometry.pt') + (folder/'camera.json').write_text(json.dumps({'distance':2, 'camera_angle_x':.7})) + Image.new('RGBA', (size,size), (30,40,50,255)).save(folder/'condition_rgba.png') + for name in ('albedo.png','mra.png'): + Image.new('RGB', (8,8), (30,40,50)).save(folder/name) + return vertices, faces + + +@pytest.mark.parametrize('with_depth', [False, True]) +def test_case_loader_keeps_geometry_bytes_and_legacy_inputs(core, monkeypatch, tmp_path, with_depth): + monkeypatch.setattr(core, 'RESOLUTION', 8) + folder = tmp_path/'case' + vertices, faces = case_files(folder, size=4) + if with_depth: + depth = np.ones((4,4), np.float32) + depth[0,0] = -1 + np.save(folder/'image_depth.npy', depth) + before = {p.name:p.read_bytes() for p in folder.iterdir()} + case = core._load_case(folder, torch, Image) + assert case['vertices'].numpy().tobytes() == vertices.numpy().tobytes() + assert case['faces'].numpy().tobytes() == faces.numpy().tobytes() + if with_depth: + assert case['image_depth'].shape == (8,8) and case['image_depth'].dtype == np.float32 + assert np.isnan(case['image_depth'][:2,:2]).all() + else: + assert case['image_depth'] is None + assert before == {p.name:p.read_bytes() for p in folder.iterdir()} + + +@pytest.mark.parametrize('with_depth', [False, True]) +def test_partial_material_observations_apply_new_gate_only_when_depth_exists(core, monkeypatch, with_depth): + monkeypatch.setattr(core, 'RESOLUTION', 4) + coords = torch.tensor([[0,0,0,2], [0,1,0,2], [0,3,3,2]], dtype=torch.int32) + source_coords = coords.clone() + surface = SimpleNamespace(coords=coords) + shape = SimpleNamespace(coords=coords) + keep = np.ones((4,4), bool); keep[0,1] = False + depth = np.ones((4,4), np.float32) + def first_hit(camera, bits, trusted, **kwargs): + assert trusted.all() + if with_depth: + assert kwargs['image_depth'] is depth + kwargs['depth_gate_out']['keep_mask'] = keep + kwargs['stats']['image_depth'] = {'status':'applied','valid_pixels':16,'rejected_pixels':1} + else: + assert 'image_depth' not in kwargs and 'depth_gate_out' not in kwargs + return object() + native = core.native_observation + monkeypatch.setattr(native, 'build_occupancy_index', lambda c: (None,None,None)) + monkeypatch.setattr(native, 'pixel_first_occupied', first_hit) + monkeypatch.setattr(native, 'owner_rows', lambda *a: torch.ones(3, dtype=torch.bool)) + encoder = Mock(return_value=object()) + monkeypatch.setattr(native, 'encode_partial', encoder) + monkeypatch.setattr(core, '_alignment_report', Mock()) + external = torch.eye(4); external[2,3] = .875 + projection = torch.tensor([[1.,0,.5],[0,1.,.5],[0,0,1.]]) + case = {'camera':{},'rgba':Image.new('RGBA',(4,4),(255,255,255,255)), + 'albedo':Image.new('RGB',(4,4),'white'),'mra':Image.new('RGB',(4,4),'white'), + 'image_depth':depth if with_depth else None} + runtime = {'torch':torch,'F':torch.nn.functional,'np':np,'Image':Image,'ImageFilter':ImageFilter,'sparse':None} + observation, metadata = core._partial(surface, shape, case, external, projection, None, {}, runtime, 'cpu') + assert observation is encoder.return_value + assert encoder.call_args.args[2].tolist() == ([True,False,True] if with_depth else [True,True,True]) + assert torch.equal(source_coords, coords) + if with_depth: + np.testing.assert_array_equal(metadata['observation_keep'], keep) + assert metadata['observation_admission']['rejected_image_depth_footprint_owners'] == 1 + else: + assert metadata == {} + + +def test_metadata_preserves_all_true_disabled_gate_and_legacy_has_no_sidecars(core, tmp_path): + output = tmp_path/'asset.npz' + core._save_observation_metadata(output, {}) + assert not list(tmp_path.iterdir()) + keep = np.ones((1024,1024), bool) + metadata = {'observation_keep':keep, 'observation_admission':{'image_depth':{ + 'status':'disabled','reason':'degenerate_image_depth_span','rejected_pixels':0}}} + core._save_observation_metadata(output, metadata) + np.testing.assert_array_equal(np.load(output.with_suffix('.observation_keep.npy'), allow_pickle=False), keep) + assert json.loads(output.with_suffix('.observation_admission.json').read_text()) == metadata['observation_admission'] + assert keep.all() + + +def test_reprojection_masks_at_1024_without_changing_source_images(core, tmp_path): + from _depth_prefilter_test import reprojection + source = tmp_path/'source'; source.mkdir() + Image.new('RGB',(8,8),(20,30,40)).save(source/'condition_rgb.png') + Image.new('RGBA',(8,8),(20,30,40,255)).save(source/'condition_rgba.png') + for name in ('albedo.png','mra.png'): + Image.new('RGB',(1024,1024),(80,90,100)).save(source/name) + camera = {'distance':2,'camera_angle_x':.7} + (source/'camera.json').write_text(json.dumps(camera)) + before = {p.name:p.read_bytes() for p in source.iterdir()} + keep = np.ones((1024,1024),bool); keep[:,501] = False + reprojection._prepare_inputs(source, tmp_path/'aligned', camera, keep=keep) + alpha = np.asarray(Image.open(tmp_path/'aligned/condition_rgba.png').getchannel('A')) + assert not alpha[:,501].any() and (alpha[:,500] == 255).all() and (alpha[:,502] == 255).all() + assert before == {p.name:p.read_bytes() for p in source.iterdir()} + + +def test_front_reprojection_forwards_gate_and_progress_without_changing_mesh(core, monkeypatch, tmp_path): + from _depth_prefilter_test import reprojection + source = tmp_path/'source'; case_files(source) + native = tmp_path/'asset.glb'; native.write_bytes(b'immutable-mesh') + keep = np.ones((1024,1024),bool); keep[100:102] = False + stage = Mock() + run_side = Mock(return_value=tmp_path/'output/front/mesh.glb') + monkeypatch.setattr(reprojection, '_run_side', run_side) + result = reprojection.run_front_view(native, source, tmp_path/'output', keep=keep, stage=stage) + assert result == run_side.return_value and native.read_bytes() == b'immutable-mesh' + assert run_side.call_args.kwargs['keep'] is keep and run_side.call_args.kwargs['stage'] is stage + with pytest.raises(ValueError, match='boolean 1024'): + reprojection.run_front_view(native, source, tmp_path/'invalid', keep=np.ones((8,8), bool)) + + +def test_public_completion_returns_metadata_and_rejects_depth_dual_view(core, monkeypatch, tmp_path): + metadata = {'observation_keep':np.ones((1024,1024),bool),'observation_admission':{}} + run = Mock(return_value=metadata) + monkeypatch.setattr(core, '_run_completion', run) + options = dict(input_dir=tmp_path, output=tmp_path/'asset.npz', model_config=tmp_path/'config', + checkpoint=tmp_path/'weights', trellis_root=tmp_path, preview_only=True, resident=object()) + assert core.run_completion(**options) is metadata + assert run.call_args.kwargs['preview_only'] and run.call_args.kwargs['resident'] is options['resident'] + np.save(tmp_path/'image_depth.npy', np.ones((2,2),np.float32)) + options.update(back_view_dir=tmp_path/'back', preview_only=False, resident=None) + with pytest.raises(ValueError, match='single-view'): + core.run_completion(**options) + assert run.call_count == 1 + + +class TinySparse: + def __init__(self, feats, coords): + self.feats, self.coords = feats, coords + self._spatial_cache = {} + + def replace(self, *, feats): + return TinySparse(feats, self.coords) + + def __mul__(self, value): + return self.replace(feats=self.feats*value) + + def __add__(self, value): + return self.replace(feats=self.feats+value) + + +@pytest.mark.parametrize('low_vram,resident_mode,preview_only,preview_failure', [ + (False, False, False, False), + (False, False, True, False), + (True, False, False, False), + (True, True, False, False), + (True, True, True, False), + (True, True, True, True), +]) +def test_both_completion_paths_save_actual_gate_and_keep_mesh_before_preview_or_export( + core, monkeypatch, tmp_path, low_vram, resident_mode, preview_only, preview_failure): + """Use tiny CPU tensors and fake model outputs to exercise actual orchestration. + + Numerical filter and interpolation tests above use their real implementations; + this test replaces only neural models/native voxelization and CUDA allocation. + """ + monkeypatch.setenv('SUPRATEX_LOW_VRAM', '1' if low_vram else '0') + original_to = torch.Tensor.to + def cpu_to(tensor, *args, **kwargs): + if args and isinstance(args[0], str) and args[0].startswith('cuda'): + args = ('cpu', *args[1:]) + if str(kwargs.get('device','')).startswith('cuda'): + kwargs['device'] = 'cpu' + return original_to(tensor, *args, **kwargs) + monkeypatch.setattr(torch.Tensor, 'to', cpu_to) + for name in ('as_tensor', 'randn'): + original = getattr(torch, name) + def allocate(*args, _original=original, **kwargs): + if str(kwargs.get('device','')).startswith('cuda'): + kwargs['device'] = 'cpu' + return _original(*args, **kwargs) + monkeypatch.setattr(torch, name, allocate) + generator = torch.Generator + monkeypatch.setattr(torch, 'Generator', lambda **kwargs: generator(device='cpu')) + monkeypatch.setattr(torch.cuda, 'is_available', lambda: True) + monkeypatch.setattr(torch.cuda, 'mem_get_info', lambda *args: (0,0)) + monkeypatch.setattr(torch.cuda, 'empty_cache', lambda: None) + + vertices = torch.tensor([[0., -.0, .25], [.5, -.5, .5], [.125, .25, -.375]]) + faces = torch.tensor([[2,0,1]], dtype=torch.int64) + original_geometry = vertices.numpy().tobytes(), faces.numpy().tobytes() + coords = torch.tensor([[0,0,0,0],[0,1,0,0],[0,0,1,0]], dtype=torch.int32) + shape = TinySparse(torch.zeros((3,32)), coords) + case = {'vertices':vertices,'faces':faces,'camera':{},'image_depth':np.ones((1024,1024),np.float32)} + keep = np.ones((1024,1024),bool); keep[400:403,100:200] = False + metadata = {'observation_keep':keep, 'observation_admission':{'image_depth':{ + 'status':'applied','rejected_pixels':300,'valid_pixels':10000}}} + partial = Mock(return_value=(SimpleNamespace(feats=torch.zeros((3,34))), metadata)) + monkeypatch.setattr(core, '_partial', partial) + monkeypatch.setattr(core, '_load_case', lambda *args: case) + monkeypatch.setattr(core, '_shape', lambda *args: (shape, shape)) + monkeypatch.setattr(core, '_camera', lambda *args: (torch.eye(4),torch.eye(3))) + monkeypatch.setattr(core, '_stage_sparse', lambda value, *args, **kwargs: value) + monkeypatch.setattr(core, '_stem', lambda path, *args: path) + monkeypatch.setattr(core, '_completion_vae_root', lambda *args: tmp_path) + normalization = {'mean':[0.]*32, 'std':[1.]*32} + dataset = {'shape_slat_normalization':normalization,'pbr_slat_normalization':normalization, + 'pretrained_shape_slat_dec':'shape_decoder', 'pretrained_pbr_slat_dec':'material_decoder'} + monkeypatch.setattr(core, 'validate_completion_artifacts', lambda *args: ({},dataset,tmp_path/'weights',{})) + monkeypatch.setattr(core, 'completion_decoder_stems', lambda *args: { + key:tmp_path/dataset[key] for key in ('pretrained_shape_slat_dec','pretrained_pbr_slat_dec')}) + monkeypatch.setattr(core, 'load_completion_state_dict', lambda *args, **kwargs: {}) + + class Model: + low_vram = False + blocks, refiners, model_channels = [], [], 1 + def __init__(self, name): self.name = name + def to(self, *args): return self + def eval(self): return self + def set_resolution(self, *args): pass + def __call__(self, value, **kwargs): + if self.name == 'shape_decoder': + return None, [value] + assert self.name == 'material_decoder' + return TinySparse(torch.zeros((3,6)), coords) + models = {name:Model(name) for name in ( + 'shape_encoder','texture_encoder','transformer','shape_decoder','material_decoder')} + def vae(_module, stem): + name = Path(stem).name + if name.startswith('shape_enc'): name = 'shape_encoder' + if name.startswith('tex_enc'): name = 'texture_encoder' + return models[name] + monkeypatch.setattr(core, '_load_vae', vae) + neural = ModuleType('supratex.pipelines.observation_guided_completion') + neural.build_completion_model = lambda *args, **kwargs: models['transformer'] + neural.sample_material_latents = lambda _model, noise, **kwargs: SimpleNamespace(samples=noise) + monkeypatch.setitem(sys.modules, neural.__name__, neural) + resident = SimpleNamespace(completion_model=Mock(side_effect=models.__getitem__)) if resident_mode else None + emitted = [] + class Mesh: + def __init__(self, vertices, faces, *args, **kwargs): + self.vertices, self.faces = vertices, faces + def to(self, device): return self + monkeypatch.setattr(core, '_runtime', lambda: { + 'torch':torch,'models':object(),'Image':Image,'utils3d':None,'MeshWithVoxel':Mesh}) + output = tmp_path/'asset.npz' + def emit(path, mesh, *args, **kwargs): + # Sidecars exist even when preview rendering later fails, for recovery. + np.testing.assert_array_equal(np.load(output.with_suffix('.observation_keep.npy'), allow_pickle=False), keep) + assert json.loads(output.with_suffix('.observation_admission.json').read_text()) == metadata['observation_admission'] + emitted.append((mesh.vertices.numpy().tobytes(), mesh.faces.numpy().tobytes())) + if preview_failure: + raise RuntimeError('preview failed after saved geometry and gate') + preview = ModuleType('hf_space.ovoxel_preview'); preview.save_and_render = emit + monkeypatch.setitem(sys.modules, preview.__name__, preview) + monkeypatch.setattr(core, '_export', emit) + options = dict(input_dir=tmp_path, output=output, model_config=tmp_path/'config', + checkpoint=tmp_path/'weights', trellis_root=tmp_path, + preview_only=preview_only, resident=resident) + if preview_failure: + with pytest.raises(RuntimeError, match='preview failed'): + core.run_completion(**options) + else: + assert core.run_completion(**options) is metadata + assert emitted == [original_geometry] + assert (vertices.numpy().tobytes(), faces.numpy().tobytes()) == original_geometry + assert partial.call_count == 1 + if resident is not None: + assert [call.args[0] for call in resident.completion_model.call_args_list] == [ + 'shape_encoder','texture_encoder','transformer','shape_decoder','material_decoder'] diff --git a/tests/test_frontend_depth.py b/tests/test_frontend_depth.py new file mode 100644 index 0000000000000000000000000000000000000000..b56389333eeb32f2903af3b4bfa5bec46855e052 --- /dev/null +++ b/tests/test_frontend_depth.py @@ -0,0 +1,92 @@ +"""Camera depth comes from the same MoGe result without changing its image grid.""" +from contextlib import nullcontext +import importlib.util +from pathlib import Path +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +from PIL import Image +import pytest + + +def camera_module(): + path = Path(__file__).resolve().parents[1] / 'supratex/frontend/camera.py' + spec = importlib.util.spec_from_file_location('_test_frontend_camera_depth', path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class Tensor: + """Small CPU adapter for the Torch conversion boundary, not model inference.""" + def __init__(self, value): + self.value = np.asarray(value) + + def detach(self): return self + def cpu(self): return self + def to(self, _device): return self + def float(self): return Tensor(self.value.astype(np.float32)) + def numpy(self): return self.value + def squeeze(self): return Tensor(self.value.squeeze()) + def permute(self, *axes): return Tensor(self.value.transpose(axes)) + + +def test_resident_camera_saves_same_inference_depth_and_valid_mask(monkeypatch, tmp_path): + camera = camera_module() + rgb = np.arange(18, dtype=np.uint8).reshape(2, 3, 3) + Image.fromarray(rgb).save(tmp_path/'input.png') + (tmp_path/'model.pt').write_bytes(b'checkpoint') + depth = np.array([[[1, 2, -1], [np.inf, 4, 5]]], dtype=np.float64) + result = {'intrinsics': Tensor(np.eye(3)), 'depth': Tensor(depth), + 'mask': Tensor(np.array([[[1, 0, 1], [1, 1, np.nan]]]))} + model = SimpleNamespace(infer=Mock(return_value=result), cpu=Mock()) + model.to = Mock(return_value=model); model.eval = Mock(return_value=model) + pretrained = Mock(side_effect=AssertionError('The resident model must be reused')) + monkeypatch.setitem(sys.modules, '_fake_moge_depth', SimpleNamespace(MoGeModel=SimpleNamespace(from_pretrained=pretrained))) + empty_cache = Mock() + monkeypatch.setitem(sys.modules, 'torch', SimpleNamespace(from_numpy=Tensor, inference_mode=nullcontext, + cuda=SimpleNamespace(empty_cache=empty_cache))) + resident = SimpleNamespace(pool=SimpleNamespace(get=Mock(return_value=model))) + output = tmp_path/'handoff'/'image_depth.npy' + value = camera.estimate(tmp_path/'input.png', tmp_path, '_fake_moge_depth', resident=resident, + depth_output=output) + resident.pool.get.assert_called_once_with('camera') + model.infer.assert_called_once() + np.testing.assert_allclose(model.infer.call_args.args[0].numpy(), rgb.astype(np.float32).transpose(2, 0, 1)/255) + assert model.infer.call_args.kwargs == {'resolution_level': 9, 'force_projection': True, 'apply_mask': True} + saved = np.load(output, allow_pickle=False) + assert saved.shape == (2, 3) and saved.dtype == np.float32 + np.testing.assert_equal(saved, [[1, np.nan, np.nan], [np.nan, 4, np.nan]]) + assert value == camera.camera_from_intrinsics(np.eye(3), 3, 2, resolution_level=9) + model.cpu.assert_called_once(); empty_cache.assert_called_once(); pretrained.assert_not_called() + with pytest.raises(FileExistsError, match='must be new'): + camera.estimate(tmp_path/'input.png', tmp_path, '_fake_moge_depth', resident=resident, depth_output=output) + assert model.infer.call_count == 1 + + +@pytest.mark.parametrize('result, message', [ + ({}, 'did not return'), + ({'depth': Tensor(np.ones((4, 4)))}, 'pixel grid'), + ({'depth': Tensor(np.ones((2, 3))), 'mask': Tensor(np.ones((2, 2)))}, 'mask must match'), +]) +def test_camera_depth_rejects_missing_or_resampled_outputs(result, message): + with pytest.raises(ValueError, match=message): + camera_module()._image_depth(result, 3, 2) + + +def test_camera_without_depth_output_preserves_legacy_call(monkeypatch, tmp_path): + camera = camera_module() + Image.new('RGB', (3, 2), 'white').save(tmp_path/'input.png') + (tmp_path/'model.pt').write_bytes(b'checkpoint') + model = SimpleNamespace(infer=Mock(return_value={'intrinsics': Tensor(np.eye(3))}), cpu=Mock()) + model.to = Mock(return_value=model); model.eval = Mock(return_value=model) + factory = Mock(return_value=model) + monkeypatch.setitem(sys.modules, '_fake_moge_legacy', SimpleNamespace(MoGeModel=SimpleNamespace(from_pretrained=factory))) + monkeypatch.setitem(sys.modules, 'torch', SimpleNamespace(from_numpy=Tensor, inference_mode=nullcontext)) + result = camera.estimate(tmp_path/'input.png', tmp_path, '_fake_moge_legacy', device='cpu') + assert result['image_width'] == 3 and result['image_height'] == 2 + factory.assert_called_once_with(str(tmp_path/'model.pt')); model.infer.assert_called_once() + model.cpu.assert_called_once() + assert not list(tmp_path.glob('*.npy')) diff --git a/tests/test_geometry_results.py b/tests/test_geometry_results.py new file mode 100644 index 0000000000000000000000000000000000000000..923d748fc4d105d1067f3a94c9ca1868d8bef4db --- /dev/null +++ b/tests/test_geometry_results.py @@ -0,0 +1,129 @@ +"""Geometry survives local cache loss and cannot resume from corrupt inputs.""" +import json +from pathlib import Path +import shutil + +import numpy as np +import pytest +from PIL import Image + +from hf_space import geometry_results, ovoxel_results, split_pipeline + + +@pytest.fixture +def geometry(tmp_path, monkeypatch): + monkeypatch.setattr(ovoxel_results, 'ROOT', tmp_path / 'durable') + folder = tmp_path / ('a' * 32) + folder.mkdir() + for name in geometry_results.HANDOFF_FILES: + path = folder / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b'fixture:' + name.encode()) + (folder / 'geometry-checkpoint' / 'checkpoint.json').write_text(json.dumps({ + 'format': split_pipeline.CHECKPOINT_FORMAT, 'seed': 42, 'remove_floaters_applied': True, + 'files': list(split_pipeline.CHECKPOINT_FILES[:-1])})) + Image.new('RGBA', (8, 8), 'white').save(folder / 'input.png') + metadata = {'views': 2, 'modes': ['clay'], 'default_mode': 'clay', + 'view_layout': 'input_front_y_up_v1', 'input_view': 0} + (folder / 'preview.json').write_text(json.dumps(metadata)) + (folder / 'report.json').write_text('{"status":"complete"}') + (folder / 'timings.jsonl').write_text('{"name":"geometry.preview"}\n') + (folder / 'views').mkdir() + for index in range(2): + Image.new('RGB', (8, 8), 'gray').save(folder / 'views' / f'clay-{index:02d}.webp') + return folder + + +def test_geometry_is_restorable_after_original_worker_files_are_gone(geometry, tmp_path): + original = {name: (geometry / name).read_bytes() for name in geometry_results.HANDOFF_FILES} + identifier = geometry_results.persist(geometry / 'geometry.npz', {'project_revision': 'test'}) + shutil.rmtree(geometry) + restored = geometry_results.restore(identifier, tmp_path / 'restored') + assert {name: (restored / name).read_bytes() for name in original} == original + assert not (restored / 'views').exists() + ready = json.loads((ovoxel_results.directory(identifier) / 'ready.json').read_text()) + assert ready['output_kind'] == 'geometry' + + +def test_modified_geometry_is_rejected_before_texture_execution(geometry, tmp_path): + identifier = geometry_results.persist(geometry / 'geometry.npz', {'project_revision': 'test'}) + (ovoxel_results.directory(identifier) / 'geometry-checkpoint' / 'geometry.pt').write_bytes(b'corrupt') + with pytest.raises(ValueError, match='integrity'): + geometry_results.restore(identifier, tmp_path / 'restored') + + +def test_partial_publish_is_not_resumable(geometry, tmp_path): + identifier = geometry_results.persist(geometry / 'geometry.npz', {'project_revision': 'test'}) + (ovoxel_results.directory(identifier) / 'ready.json').unlink() + with pytest.raises(ValueError, match='incomplete'): + geometry_results.restore(identifier, tmp_path / 'restored') + + +def test_unexpected_checkpoint_paths_are_rejected(geometry, tmp_path): + identifier = geometry_results.persist(geometry / 'geometry.npz', {'project_revision': 'test'}) + path = ovoxel_results.directory(identifier) / 'geometry-manifest.json' + manifest = json.loads(path.read_text()) + manifest['files']['../../outside'] = {'size': 1, 'sha256': '0' * 64} + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match='handoff'): + geometry_results.restore(identifier, tmp_path / 'restored') + assert not (tmp_path / 'restored').exists() + + +def test_incomplete_white_preview_is_not_published(geometry): + (geometry / 'views' / 'clay-01.webp').unlink() + with pytest.raises(ValueError, match='preview is incomplete'): + geometry_results.persist(geometry / 'geometry.npz', {'project_revision': 'test'}) + assert not ovoxel_results.directory(geometry.name).exists() + + +def add_depth(folder): + checkpoint = folder/'geometry-checkpoint' + np.save(checkpoint/'image_depth.npy', np.array([[1, np.nan], [2, 3]], dtype=np.float32), allow_pickle=False) + metadata = json.loads((checkpoint/'checkpoint.json').read_text()) + metadata['files'].append('image_depth.npy') + (checkpoint/'checkpoint.json').write_text(json.dumps(metadata)) + return (checkpoint/'image_depth.npy').read_bytes() + + +def test_depth_is_bound_and_restored_without_changing_saved_geometry(geometry, tmp_path): + depth = add_depth(geometry) + mesh = (geometry/'geometry.npz').read_bytes() + checkpoint_mesh = (geometry/'geometry-checkpoint'/'geometry.pt').read_bytes() + identifier = geometry_results.persist(geometry/'geometry.npz', {'project_revision': 'test'}) + durable = ovoxel_results.directory(identifier) + manifest = json.loads((durable/'geometry-manifest.json').read_text()) + assert set(manifest['files']) == {*geometry_results.HANDOFF_FILES, 'geometry-checkpoint/image_depth.npy'} + assert manifest['files']['geometry-checkpoint/image_depth.npy']['sha256'] == geometry_results._digest(durable/'geometry-checkpoint'/'image_depth.npy') + shutil.rmtree(geometry) + restored = geometry_results.restore(identifier, tmp_path/'restored') + assert (restored/'geometry-checkpoint'/'image_depth.npy').read_bytes() == depth + assert (restored/'geometry.npz').read_bytes() == mesh + assert (restored/'geometry-checkpoint'/'geometry.pt').read_bytes() == checkpoint_mesh + + +@pytest.mark.parametrize('damage, message', [('missing', 'incomplete'), ('changed', 'integrity'), ('unbound', 'handoff')]) +def test_new_depth_cannot_be_lost_or_changed_during_restore(geometry, tmp_path, damage, message): + add_depth(geometry) + identifier = geometry_results.persist(geometry/'geometry.npz', {'project_revision': 'test'}) + durable = ovoxel_results.directory(identifier) + depth = durable/'geometry-checkpoint'/'image_depth.npy' + if damage == 'missing': + depth.unlink() + elif damage == 'changed': + depth.write_bytes(b'corrupt depth') + else: + path = durable/'geometry-manifest.json' + manifest = json.loads(path.read_text()) + del manifest['files']['geometry-checkpoint/image_depth.npy'] + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match=message): + geometry_results.restore(identifier, tmp_path/'restored') + + +def test_declared_missing_depth_is_not_published(geometry): + add_depth(geometry) + (geometry/'geometry-checkpoint'/'image_depth.npy').unlink() + with pytest.raises(ValueError, match='incomplete'): + geometry_results.persist(geometry/'geometry.npz', {'project_revision': 'test'}) + assert not ovoxel_results.directory(geometry.name).exists() diff --git a/tests/test_history_geometry_texture.py b/tests/test_history_geometry_texture.py new file mode 100644 index 0000000000000000000000000000000000000000..fad8e543e32ca498a162f91b1218703968f490f9 --- /dev/null +++ b/tests/test_history_geometry_texture.py @@ -0,0 +1,344 @@ +"""The durable geometry checkpoint survives a separately requested texture stage.""" +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from io import BytesIO +import json +import time +from types import SimpleNamespace +import uuid + +from fastapi import FastAPI +from PIL import Image +import pytest +from starlette.testclient import TestClient + +from hf_space.history_auth import Identity, csrf_token +from hf_space.history_bridge import AppSessions, AppSessionMiddleware, COOKIE_NAME +from hf_space.history_jobs import Runner +from hf_space.history_routes import install_routes +from hf_space.history_store import Conflict, LimitExceeded, NotFound, Store + + +@pytest.fixture +def store(tmp_path, monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + for name in ('SPACE_ID', 'SPACE_HOST', 'OAUTH_CLIENT_ID', 'OAUTH_CLIENT_SECRET'): + monkeypatch.delenv(name, raising=False) + return Store(tmp_path/'history', asset_root=tmp_path/'assets').initialize() + + +def submit(store, owner='alice', key=None): + stream = BytesIO() + Image.new('RGB', (12, 12), 'red').save(stream, format='PNG') + upload = store.save_upload(owner, stream.getvalue()) + return store.submit(owner, upload['id'], 42, key or uuid.uuid4().hex) + + +def asset(store, mode='clay'): + identifier = uuid.uuid4().hex + folder = store.asset_root/identifier + (folder/'views').mkdir(parents=True) + Image.new('RGB', (12, 12), 'gray').save(folder/'views'/f'{mode}-00.webp') + metadata = {'modes': [mode], 'views': 1, 'default_mode': mode, 'resolution': 12} + (folder/'preview.json').write_text(json.dumps(metadata)) + (folder/'ready.json').write_text('{}') + return {'asset_id': identifier, 'preview_metadata': metadata, + 'file_manifest': {f'frames/{mode}-00.webp': { + 'size': (folder/'views'/f'{mode}-00.webp').stat().st_size, + 'version': uuid.uuid4().hex}}} + + +def geometry(store, owner='alice'): + job = submit(store, owner) + operation = job['operations'][0]['id'] + store.claim_operation(operation) + return store.finish_operation(operation, asset(store)) + + +def texture_op(store, job, key=None, owner='alice'): + requested = store.submit_texture(owner, job['id'], key or uuid.uuid4().hex) + return requested['operations'][-1]['id'] + + +def test_geometry_is_complete_viewable_and_export_requires_explicit_texture(store): + initial = submit(store) + assert initial['workflow'] == 'geometry_texture' + assert initial['texture_status'] == 'not_requested' + with pytest.raises(Conflict): + texture_op(store, initial) + operation = initial['operations'][0]['id'] + store.claim_operation(operation) + completed = store.finish_operation(operation, asset(store)) + assert completed['geometry_asset_id'] == completed['asset_id'] + assert completed['status'] == 'complete' + assert completed['texture_status'] == 'not_requested' + assert store.next_operation() is None + assert store.get_preview('alice', initial['id'])['modes'] == ['clay'] + assert store.resolve_file('alice', initial['id'], 'frames/clay-00.webp').is_file() + with pytest.raises(Conflict, match='texture'): + store.submit_export('alice', initial['id'], 'blocked-export') + + +def test_texture_deduplicates_concurrent_clicks_and_scopes_keys_to_job_owner(store): + job, other = geometry(store), geometry(store) + keys = ['texture-click-one', 'texture-click-two'] + with ThreadPoolExecutor(max_workers=2) as pool: + result = list(pool.map(lambda key: store.submit_texture('alice', job['id'], key), keys)) + assert result[0]['operations'][-1]['id'] == result[1]['operations'][-1]['id'] + current = store.get_job('alice', job['id']) + assert len([op for op in current['operations'] if op['kind'] == 'texture']) == 1 + public = store.public_snapshot(current) + assert public['status'] == 'complete' + assert public['texture_status'] == 'queued' + assert public['gpu_operation_kind'] == 'texture' + assert public['gpu_operation_id'] == current['operations'][-1]['id'] + assert 'owner' not in public and 'geometry_asset_id' not in public + with pytest.raises(Conflict): + store.submit_texture('alice', other['id'], keys[0]) + with pytest.raises(NotFound): + store.submit_texture('bob', job['id'], keys[0]) + restored = Store(store.root, asset_root=store.asset_root).initialize() + assert restored.submit_texture('alice', job['id'], keys[1])['operations'] == current['operations'] + + +@pytest.mark.parametrize('terminal', ['failed', 'interrupted']) +def test_failed_or_interrupted_texture_keeps_geometry_and_can_retry_only_texture(store, terminal): + job = geometry(store) + original_files = deepcopy(job['files']) + + def assert_single_case(current_store): + items = current_store.list_jobs('alice')['items'] + assert len(items) == 1 + assert items[0]['id'] == job['id'] + assert items[0]['created_at'] == job['created_at'] + + assert_single_case(store) + operation = texture_op(store, job, 'texture-attempt-one') + store.claim_operation(operation) + store.update_progress(operation, {'label': 'Texturing', 'elapsed_seconds': 3}) + assert_single_case(store) + assert store.get_job('alice', job['id'])['progress']['label'] == 'Complete' + if terminal == 'failed': + store.finish_operation(operation, error='Texture failed') + restored = Store(store.root, asset_root=store.asset_root).initialize() + current = restored.get_job('alice', job['id']) + assert_single_case(restored) + assert current['status'] == 'complete' + assert current['texture_status'] == terminal + assert current['asset_id'] == current['geometry_asset_id'] == job['asset_id'] + assert current['files'] == original_files + assert restored.resolve_file('alice', job['id'], 'frames/clay-00.webp').is_file() + with pytest.raises(Conflict): + restored.submit_export('alice', job['id'], 'blocked-export') + # A replay is observational; a new explicit attempt gets a new operation. + assert restored.submit_texture('alice', job['id'], 'texture-attempt-one')['texture_status'] == terminal + retry = restored.submit_texture('alice', job['id'], 'texture-attempt-two') + assert retry['operations'][-1]['id'] != operation + assert retry['texture_status'] == 'queued' and retry['texture_error'] is None + assert len([op for op in retry['operations'] if op['kind'] == 'generate']) == 1 + assert_single_case(restored) + retry_operation = retry['operations'][-1]['id'] + restored.claim_operation(retry_operation) + textured = asset(restored, 'base_color') + completed = restored.finish_operation(retry_operation, textured) + assert completed['texture_status'] == 'complete' + assert completed['asset_id'] == textured['asset_id'] != job['asset_id'] + assert completed['geometry_asset_id'] == job['asset_id'] + assert completed['preview_metadata']['modes'] == ['base_color'] + assert (restored.asset_root/job['asset_id']/'views'/'clay-00.webp').is_file() + assert_single_case(restored) + assert_single_case(Store(store.root, asset_root=store.asset_root).initialize()) + + +def test_texture_completion_atomically_switches_assets_and_keeps_input(store): + job = geometry(store) + source = store.resolve_file('alice', job['id'], 'input.png').read_bytes() + original_input_metadata = deepcopy(job['files']['input.png']) + operation = texture_op(store, job) + store.claim_operation(operation) + textured = asset(store, 'basecolor') + current = store.finish_operation(operation, textured) + assert current['status'] == current['texture_status'] == 'complete' + assert current['asset_id'] == textured['asset_id'] != current['geometry_asset_id'] + assert current['geometry_asset_id'] == job['asset_id'] + assert current['preview_metadata']['modes'] == ['basecolor'] + assert 'frames/clay-00.webp' not in current['files'] + assert current['files']['input.png'] == original_input_metadata + assert store.resolve_file('alice', job['id'], 'input.png').read_bytes() == source + descriptor = store.get_file_descriptor('alice', job['id'], 'frames/basecolor-00.webp') + assert textured['asset_id'] in descriptor['path'] + with pytest.raises(NotFound): + store.get_file_descriptor('alice', job['id'], 'frames/clay-00.webp') + assert store.submit_texture('alice', job['id'], 'another-complete-click')['operations'][-1]['id'] == operation + exported = store.submit_export('alice', job['id'], 'export-after-texture') + export_operation = exported['operations'][-1]['id'] + store.claim_operation(export_operation) + (store.asset_root/textured['asset_id']/'asset.glb').write_bytes(b'glTFtest') + current = store.finish_operation(export_operation) + assert current['export_status'] == 'complete' + assert store.resolve_file('alice', job['id'], 'asset.glb').read_bytes() == b'glTFtest' + + +@pytest.mark.parametrize('bad_result', ['same_asset', 'replace_input', 'unready_asset', 'missing_frames']) +def test_invalid_texture_result_fails_without_replacing_geometry(store, bad_result): + job = geometry(store) + operation = texture_op(store, job) + result = asset(store, 'basecolor') + if bad_result == 'same_asset': + result['asset_id'] = job['geometry_asset_id'] + elif bad_result == 'replace_input': + result['file_manifest']['input.png'] = job['files']['input.png'] + elif bad_result == 'unready_asset': + (store.asset_root/result['asset_id']/'ready.json').unlink() + else: + result['file_manifest'] = {} + calls = [] + def handle(current, update): + calls.append(current['geometry_asset_id']) + return result + runner = Runner(store, {'texture': handle}) + assert runner.run_once(operation_id=operation, owner='alice') + assert not runner.run_once(operation_id=operation, owner='alice') + current = store.get_job('alice', job['id']) + assert calls == [job['geometry_asset_id']] + assert current['status'] == 'complete' and current['texture_status'] == 'failed' + assert current['asset_id'] == job['asset_id'] + assert current['files'] == job['files'] + + +def test_pending_texture_occupies_admission_slot_and_blocks_recovery_and_deletion(store): + job = geometry(store) + operation = texture_op(store, job) + store.max_pending_per_owner = 1 + assert store.has_pending() + with pytest.raises(LimitExceeded): + submit(store) + with pytest.raises(Conflict): + store.submit_recover('alice', job['id'], 'recover-overlap') + with pytest.raises(Conflict): + store.delete_job('alice', job['id']) + store.claim_operation(operation) + store.finish_operation(operation, error='Cancelled') + assert not store.has_pending() + store.delete_job('alice', job['id']) + assert (store.asset_root/job['geometry_asset_id']/'ready.json').is_file() + with pytest.raises(NotFound): + store.get_job('alice', job['id']) + + +def test_texture_terminal_write_failure_retries_save_without_rerunning_gpu(store, monkeypatch): + job = geometry(store) + operation = texture_op(store, job) + result = asset(store, 'basecolor') + calls = [] + write = store._write_json + failed = False + + def flaky_write(path, value): + nonlocal failed + if value.get('texture_status') == 'complete' and not failed: + failed = True + raise OSError('temporary storage error') + return write(path, value) + + def handle(current, update): + calls.append(current['geometry_asset_id']) + return result + + monkeypatch.setattr(store, '_write_json', flaky_write) + runner = Runner(store, {'texture': handle}) + with pytest.raises(OSError): + runner.run_once(operation_id=operation, owner='alice') + current = store.get_job('alice', job['id']) + assert current['status'] == 'complete' and current['texture_status'] == 'running' + assert current['asset_id'] == job['asset_id'] + assert runner.run_once(operation_id=operation, owner='alice') + assert calls == [job['geometry_asset_id']] + current = store.get_job('alice', job['id']) + assert current['texture_status'] == 'complete' and current['asset_id'] == result['asset_id'] + # The existing removal policy is reversible, so it retains both asset trees. + store.delete_job('alice', job['id']) + assert (store.asset_root/current['asset_id']/'ready.json').is_file() + assert (store.asset_root/current['geometry_asset_id']/'ready.json').is_file() + + +def test_texture_admission_respects_other_queued_generation(store): + job = geometry(store) + store.max_pending_per_owner = 1 + pending = submit(store) + with pytest.raises(LimitExceeded): + texture_op(store, job) + assert store.get_job('alice', job['id'])['texture_status'] == 'not_requested' + op = pending['operations'][0]['id'] + store.claim_operation(op) + store.finish_operation(op, error='Ended') + texture_op(store, job) + assert store.get_job('alice', job['id'])['texture_status'] == 'queued' + + +def test_legacy_record_stays_exportable_after_restart(store): + job = geometry(store) + path = store.root/'jobs'/job['id']/'job.json' + value = json.loads(path.read_text()) + for name in ('workflow', 'geometry_asset_id', 'texture_status', 'texture_progress', 'texture_error'): + value.pop(name) + path.write_text(json.dumps(value)) + restored = Store(store.root, asset_root=store.asset_root).initialize() + assert not restored.get_job('alice', job['id']).get('workflow') + with pytest.raises(Conflict): + restored.submit_texture('alice', job['id'], 'legacy-no-texture') + assert restored.submit_export('alice', job['id'], 'legacy-export')['export_status'] == 'queued' + + +@pytest.fixture +def web(store, tmp_path): + sessions = AppSessions('test-only-texture-session') + app = FastAPI() + install_routes(app, store, deployment_status=lambda: {'gpu_mode': 'zerogpu'}, + private_cache_root=tmp_path/'cache') + app.add_middleware(AppSessionMiddleware, sessions=sessions) + with TestClient(app, base_url='https://testserver') as client: + yield SimpleNamespace(client=client, sessions=sessions, store=store) + + +def headers(web, owner='alice'): + claims = web.sessions.claims(Identity(owner, owner, 'test-oauth'), time.time()+600) + identity = web.sessions.validate(claims) + return {'Origin': 'https://testserver', 'X-CSRF-Token': csrf_token(identity), + 'Cookie': COOKIE_NAME+'='+web.sessions.serializer.dumps(claims)} + + +def test_texture_route_auth_csrf_schema_and_idempotency(web): + job = geometry(web.store) + path = f'/api/history/jobs/{job["id"]}/texture' + body = {'idempotency_key': 'route-texture-click'} + assert web.client.post(path, json=body).status_code == 401 + assert web.client.post(path, headers=headers(web, 'bob'), json=body).status_code == 404 + missing_csrf = headers(web) + missing_csrf.pop('X-CSRF-Token') + assert web.client.post(path, headers=missing_csrf, json=body).status_code == 403 + assert web.client.post(path, headers=headers(web), json={**body, 'asset_id': job['asset_id']}).status_code == 400 + request_headers = headers(web) + response = web.client.post(path, headers=request_headers, json=body) + assert response.status_code == 200, response.text + value = response.json() + assert value['workflow'] == 'geometry_texture' and value['texture_status'] == 'queued' + assert value['status'] == 'complete' and value['gpu_operation_kind'] == 'texture' + replay = web.client.post(path, headers=request_headers, json=body) + assert replay.json()['gpu_operation_id'] == value['gpu_operation_id'] + frame = web.client.get(f'/api/history/jobs/{job["id"]}/files/frames/clay-00.webp', headers=request_headers) + assert frame.status_code == 200 and frame.headers['content-type'] == 'image/webp' + blocked = web.client.post(f'/api/history/jobs/{job["id"]}/export', headers=request_headers, + json={'idempotency_key': 'export-too-early'}) + assert blocked.status_code == 409 + + +def test_submit_route_marks_new_workflow(web): + stream = BytesIO() + Image.new('RGB', (12, 12), 'red').save(stream, format='PNG') + upload = web.store.save_upload('alice', stream.getvalue()) + response = web.client.post('/api/history/submit', headers=headers(web), json={ + 'upload_id': upload['id'], 'seed': 42, 'idempotency_key': 'route-first-geometry'}) + assert response.status_code == 200, response.text + assert response.json()['workflow'] == 'geometry_texture' + assert response.json()['texture_status'] == 'not_requested' diff --git a/tests/test_history_service_workflow.py b/tests/test_history_service_workflow.py new file mode 100644 index 0000000000000000000000000000000000000000..bcebc81ad7d5ffce84169b873b4be57077532d35 --- /dev/null +++ b/tests/test_history_service_workflow.py @@ -0,0 +1,155 @@ +"""Exercise public split jobs and operator full jobs through their real services.""" +import hashlib +from io import BytesIO +import json +from pathlib import Path +from types import SimpleNamespace +import uuid + +from PIL import Image +import pytest + +from hf_space import generation_progress, geometry_results, history_service, ovoxel_results, split_pipeline +from hf_space.curated_refresh import Coordinator, PLAN +from hf_space.history_store import Store + + +@pytest.fixture +def service(tmp_path, monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + for name in ('SPACE_ID', 'SPACE_HOST'): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(ovoxel_results, 'ROOT', tmp_path/'durable') + monkeypatch.setattr(ovoxel_results, 'JOBS', tmp_path/'worker') + monkeypatch.setattr(generation_progress, 'DIRECTORY', tmp_path/'progress') + monkeypatch.setattr(generation_progress, 'JOBS', tmp_path/'worker') + monkeypatch.setattr(history_service.Service, 'ready', staticmethod(lambda kind: True)) + store = Store(tmp_path/'history', asset_root=ovoxel_results.ROOT).initialize() + result = history_service.Service(store=store) + result.private_files = SimpleNamespace(seed=lambda *args, **kwargs: None) + return result + + +def generated_output(*, geometry): + folder = ovoxel_results.JOBS/uuid.uuid4().hex + folder.mkdir(parents=True) + mode = 'clay' if geometry else 'base_color' + filename = 'geometry.npz' if geometry else 'asset.npz' + Image.new('RGBA', (8, 8), 'red').save(folder/'input.png') + (folder/filename).write_bytes(b'saved-geometry' if geometry else b'saved-texture') + if geometry: + for name in geometry_results.HANDOFF_FILES: + path = folder/name + if not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b'checkpoint:' + name.encode()) + (folder/'geometry-checkpoint'/'checkpoint.json').write_text(json.dumps({ + 'format': split_pipeline.CHECKPOINT_FORMAT, 'seed': 42, 'remove_floaters_applied': True, + 'files': list(split_pipeline.CHECKPOINT_FILES[:-1])})) + report = {'status': 'complete', 'project_revision': 'fixture', + 'output_kind': 'geometry' if geometry else 'final_ovoxel'} + (folder/'report.json').write_text(json.dumps(report)) + (folder/'timings.jsonl').write_text('{"name":"fixture"}\n') + (folder/'preview.json').write_text(json.dumps({ + 'modes': [mode], 'views': 1, 'default_mode': mode, 'resolution': 8, + 'view_layout': 'input_front_y_up_v1', 'input_view': 0})) + (folder/'views').mkdir() + Image.new('RGB', (8, 8), 'gray').save(folder/'views'/f'{mode}-00.webp') + return str(folder/filename), str(folder/filename), report + + +def new_job(service): + data = BytesIO() + Image.new('RGBA', (8, 8), 'red').save(data, format='PNG') + upload = service.store.save_upload('alice', data.getvalue()) + return service.store.submit('alice', upload['id'], 42, uuid.uuid4().hex) + + +def test_service_runs_geometry_then_restores_exact_checkpoint_only_for_explicit_texture(service, monkeypatch): + calls, restored_folders = [], [] + + def generate_geometry(image, seed, **kwargs): + calls.append(('geometry', seed)) + return generated_output(geometry=True) + + def generate_texture(folder, seed, **kwargs): + restored_folders.append(Path(folder)) + assert (folder/'geometry-checkpoint'/'geometry.pt').read_bytes() == b'checkpoint:geometry-checkpoint/geometry.pt' + assert (folder/'geometry.npz').read_bytes() == b'saved-geometry' + assert all((folder/name).is_file() for name in geometry_results.HANDOFF_FILES) + calls.append(('texture', seed)) + return generated_output(geometry=False) + + monkeypatch.setattr(history_service.release_backend, 'generate_geometry', generate_geometry) + monkeypatch.setattr(history_service.release_backend, 'generate_texture', generate_texture) + monkeypatch.setattr(history_service.release_backend, 'generate', lambda *a, **kw: pytest.fail('No full inference for split job')) + job = new_job(service) + assert service.runner.run_once(operation_id=job['operations'][0]['id'], owner='alice') + white = service.store.get_job('alice', job['id']) + assert white['status'] == 'complete' and white['texture_status'] == 'not_requested' + assert calls == [('geometry', 42)] + assert service.store.next_operation() is None + queued = service.store.submit_texture('alice', job['id'], 'explicit-texture-click') + assert service.runner.run_once(operation_id=queued['operations'][-1]['id'], owner='alice') + textured = service.store.get_job('alice', job['id']) + assert calls == [('geometry', 42), ('texture', 42)] + assert textured['status'] == textured['texture_status'] == 'complete' + assert textured['geometry_asset_id'] == white['asset_id'] != textured['asset_id'] + assert textured['preview_metadata']['modes'] == ['base_color'] + assert all(not path.exists() for path in restored_folders) + assert (ovoxel_results.directory(textured['asset_id'])/'asset.npz').read_bytes() == b'saved-texture' + + +def test_corrupt_saved_geometry_fails_before_texture_gpu_callback(service, monkeypatch): + monkeypatch.setattr(history_service.release_backend, 'generate_geometry', + lambda *a, **kw: generated_output(geometry=True)) + monkeypatch.setattr(history_service.release_backend, 'generate_texture', + lambda *a, **kw: pytest.fail('Invalid checkpoint must not acquire GPU')) + job = new_job(service) + service.runner.run_once(operation_id=job['operations'][0]['id'], owner='alice') + white = service.store.get_job('alice', job['id']) + (ovoxel_results.directory(white['geometry_asset_id'])/'geometry-checkpoint'/'geometry.pt').write_bytes(b'corrupt') + queued = service.store.submit_texture('alice', job['id'], 'corrupt-texture-click') + service.runner.run_once(operation_id=queued['operations'][-1]['id'], owner='alice') + current = service.store.get_job('alice', job['id']) + assert current['status'] == 'complete' and current['texture_status'] == 'failed' + assert current['asset_id'] == white['asset_id'] + assert service.store.resolve_file('alice', job['id'], 'frames/clay-00.webp').is_file() + + +def test_curated_refresh_keeps_full_pipeline_and_can_queue_export(service, tmp_path, monkeypatch): + monkeypatch.delenv('SPACES_ZERO_GPU') + image = tmp_path/'approved.png' + rgba = Image.new('RGBA', (8, 8), 'red') + rgba.save(image) + pixel_hash = hashlib.sha256(f'RGBA:{rgba.width}:{rgba.height}:'.encode() + rgba.tobytes()).hexdigest() + plan = json.loads(PLAN.read_text()) + plan['enabled'] = True + for entry in plan['examples']: + entry.update(input_sha256=hashlib.sha256(image.read_bytes()).hexdigest(), input_rgba_sha256=pixel_hash) + monkeypatch.setattr(service, 'example_file', lambda name, kind: image) + calls = [] + + def full_generation(image, seed, **kwargs): + calls.append(seed) + return generated_output(geometry=False) + + monkeypatch.setattr(history_service.release_backend, 'generate', full_generation) + monkeypatch.setattr(history_service.release_backend, 'generate_geometry', + lambda *a, **kw: pytest.fail('Curated refresh must use the full pipeline')) + coordinator = Coordinator(service, plan=plan, state_root=tmp_path/'refresh') + assert coordinator.run_once() + item = coordinator.state['items'][plan['examples'][0]['name']] + job = service.store.get_job(coordinator.owner, item['job_id']) + assert job.get('workflow') is None + assert service.runner.run_once() + assert calls == [42] + assert coordinator.run_once() + current = service.store.get_job(coordinator.owner, item['job_id']) + assert current['status'] == 'complete' and current['export_status'] == 'queued' + assert [op['kind'] for op in current['operations']] == ['generate', 'export'] + assert item['stage'] == 'export_queued' + # Admission replay after a coordinator restart cannot create another job. + restored = Coordinator(service, plan=plan, state_root=tmp_path/'refresh') + assert not restored.run_once() + assert len(service.store.list_jobs(coordinator.owner)['items']) == 1 diff --git a/tests/test_history_ui_resume.cjs b/tests/test_history_ui_resume.cjs new file mode 100644 index 0000000000000000000000000000000000000000..6bf87bba8a392cd554d784e46a12b4e2a2d813e9 --- /dev/null +++ b/tests/test_history_ui_resume.cjs @@ -0,0 +1,811 @@ +const test = require('node:test'); +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const path = require('node:path'); +const vm = require('node:vm'); + +const read = name => fs.readFileSync(path.join(__dirname, '../hf_space', name), 'utf8'); +const uiSource = read('history_ui.js'); +const markup = read('history_ui.py').match(/MARKUP = """([\s\S]*?)"""/)[1]; +const tick = () => new Promise(resolve => setImmediate(resolve)); +const deferred = () => { let resolve; const promise = new Promise(r => { resolve = r; }); return {promise, resolve}; }; +const operationId = 'a'.repeat(32); +const savedJob = {id: 'saved-job', status: 'queued', gpu_mode: 'zerogpu', + gpu_operation_id: operationId, gpu_operation_kind: 'generate', progress: {}}; +const deployment = (runtimeId, completed = 2, total = 5) => ({ + gpu_mode: 'zerogpu', runtime_id: runtimeId, workspace: {stage: 'ready'}, + native_build: {stage: 'compiled'}, foreground: {stage: 'pending'}, + models: {stage: 'loading_models', startup_progress: {completed, total, unit: 'models', label: 'Internal model name'}} +}); + +function fakeDOM() { + const nodes = new Map(); + class Element { + constructor(tagName = 'div') { + this.tagName = tagName; this.dataset = {}; this.style = {}; this.attributes = new Map(); + this.children = []; this.handlers = new Map(); this.textHistory = []; this.textContent = ''; this.hidden = false; + this.disabled = false; this.inert = false; this.offsetHeight = 44; + this.classList = {toggle() {}, add() {}, remove() {}}; + } + get textContent() { return this._textContent; } + set textContent(value) { this._textContent = value; this.textHistory.push(value); } + close() { this.hidden = true; } + addEventListener(type, handler) { + const handlers = this.handlers.get(type) || []; + handlers.push(handler); this.handlers.set(type, handlers); + } + click() { return Promise.all((this.handlers.get('click') || []).map(handler => handler({target: this}))); } + append(...children) { + children = children.flatMap(child => child.tagName === '#fragment' ? child.children : [child]); + this.children.push(...children); + for (const child of children) if (child.id) nodes.set(child.id, child); + } + replaceChildren(...children) { this.children = []; this.append(...children); } + querySelector(selector) { + if (selector.startsWith('#')) return nodes.get(selector.slice(1)) || null; + return this.querySelectorAll(selector)[0] || null; + } + querySelectorAll(selector) { + const match = node => selector.startsWith('.') ? (node.className || '').split(' ').includes(selector.slice(1)) : + selector.includes('[') ? node.tagName === selector.split('[')[0] && + selector.match(/\[data-([\w-]+)/) && selector.match(/\[data-([\w-]+)/)[1].replace(/-([a-z])/g, (_, value) => value.toUpperCase()) in node.dataset : + node.tagName === selector; + return this.children.flatMap(child => [...(match(child) ? [child] : []), ...child.querySelectorAll(selector)]); + } + contains(node) { return node === this || this.children.some(child => child.contains(node)); } + setAttribute(name, value) { this.attributes.set(name, String(value)); } + getAttribute(name) { return this.attributes.get(name) || null; } + removeAttribute(name) { this.attributes.delete(name); if (name === 'value') delete this.value; } + } + // Derive available IDs and initial hidden states from the actual page markup. + // Missing/removed elements return null, so stale production DOM accesses fail. + for (const tag of markup.matchAll(/<([a-z][\w-]*)\b([^>]*)>/gi)) { + const id = tag[2].match(/\bid="([^"]+)"/); + if (!id) continue; + const node = new Element(tag[1]); node.id = id[1]; + node.hidden = /\bhidden\b/.test(tag[2]); nodes.set(node.id, node); + } + assert.match(markup, /id="st-job-progress"[\s\S]*? nodes.get(id) || null, handlers: documentHandlers, + createElement: tag => new Element(tag), createDocumentFragment: () => new Element('#fragment'), + addEventListener(type, handler) { + const handlers = documentHandlers.get(type) || []; + handlers.push(handler); documentHandlers.set(type, handlers); + }}; + return {document, get: id => nodes.get(id) || null}; +} + +function setup(t) { + const dom = fakeDOM(), requests = [], submissions = [], finishGPU = deferred(); + const network = {deployment: deployment('runtime-a'), statusFailure: false, nextStatus: null}; + const images = {decoded: 0, created: [], revoked: []}; + class UploadURL extends URL { + static createObjectURL(blob) { const url = 'blob:test-' + (images.created.length + 1); images.created.push({url, blob}); return url; } + static revokeObjectURL(url) { images.revoked.push(url); } + } + class UploadImage { + async decode() { images.decoded++; network.onDecode?.(); if (network.nextDecode) { const pending = network.nextDecode; network.nextDecode = null; await pending.promise; } } + } + let timerId = 0; + const timers = new Map(); + const response = (data, status = 200) => new Response(JSON.stringify(data), {status}); + const context = vm.createContext({ + AbortController, AbortSignal, DOMException, Headers, URL: UploadURL, URLSearchParams, Image: UploadImage, crypto: require('node:crypto').webcrypto, + document: dom.document, window: {addEventListener() {}}, + location: {origin: 'https://example.hf.space', href: 'https://example.hf.space/'}, + matchMedia: () => ({matches: false}), + requestAnimationFrame: callback => callback(), + setTimeout(callback) { const id = ++timerId; timers.set(id, callback); return id; }, + clearTimeout: id => timers.delete(id), + async fetch(url, options = {}) { + requests.push({url, method: options.method || 'GET', headers: options.headers, body: options.body}); + if (url === '/api/history/status') { + if (network.nextStatus) { + const pending = network.nextStatus; network.nextStatus = null; return pending.promise; + } + return network.statusFailure ? response({detail: 'Temporarily unavailable'}, 503) : response(network.deployment); + } + if (/\/uploads\/[^/]+\/files\/foreground\.png$/.test(url)) { + if (network.nextImage) { const pending = network.nextImage; network.nextImage = null; return pending.promise; } + return network.imageFailure ? new Response('', {status: 503}) : new Response(new Blob(['cutout'], {type: 'image/png'})); + } + if (/^\/api\/history\/uploads\/[^/]+$/.test(url)) { + if (network.nextUploadStatus) { const pending = network.nextUploadStatus; network.nextUploadStatus = null; return pending.promise; } + return response(network.uploadResponse); + } + if (/^\/api\/history\/uploads\/[^/]+\/foreground$/.test(url)) return response(network.retryUploadResponse); + if (url === '/api/history/gpu/prepare') return response({}); + if (url.startsWith('/api/history/jobs?') && network.jobsResponse) return response(network.jobsResponse); + if (url === '/api/history/submit' && network.nextSubmission) { + const pending = network.nextSubmission; network.nextSubmission = null; return pending.promise; + } + if (url === '/api/history/jobs/saved-job/texture') return response(network.textureResponse, network.textureFailure ? 503 : 200); + const previewJob = /^\/api\/history\/jobs\/([^/]+)\/preview\?/.exec(url); + if (previewJob) { + if (network.nextPreview) { + const pending = network.nextPreview; network.nextPreview = null; return pending.promise; + } + return response({view: 'render', html: network.previewByJob?.[previewJob[1]] || network.previewHTML || '
White preview
'}); + } + throw new Error('Unexpected request in Resume integration test: ' + url); + } + }); + vm.runInContext(read('zerogpu_client.js'), context); + const realClient = context.SupraZeroGPU; + context.window.SupraZeroGPU = {...realClient, createDispatcher: options => realClient.createDispatcher({ + ...options, + quotaHeaders: async () => ({'x-zerogpu-token': 'test-only-token'}), + connect: async () => ({submit(endpoint, data) { + submissions.push({endpoint, data}); + return (async function* () { await finishGPU.promise; })(); + }}) + })}; + assert.equal((uiSource.match(/void boot\(\);/g) || []).length, 1); + vm.runInContext(uiSource.replace('void boot();', + 'globalThis.testUI = {S, renderSelected, refreshStatus, refreshJobs, updateGate, submit, gpuDispatcher, resumeGPU, selectJob, selectExample, previewKey, historyStatus, applyUploadResult, refreshUpload, clearPrivate, clearUploadPreview, retryUploadPreview};'), context); + const ui = context.testUI; + Object.assign(ui.S, {identity: {authenticated: true, csrf_token: 'csrf-test'}, authChecking: false, + selected: {kind: 'job', id: savedJob.id}, source: {upload_id: 'saved-upload'}, + historyLoaded: true, examplesLoaded: true, + deployment: {...deployment('previous-runtime'), models: {stage: 'ready_released'}}}); + ui.S.jobs.set(savedJob.id, {...savedJob}); + ui.renderSelected(); + t.after(async () => { + ui.S.stopped = true; ui.gpuDispatcher.reset(); finishGPU.resolve(); await tick(); + }); + return {...ui, ...dom, network, requests, submissions, response, images, timers, + prepareRequests: () => requests.filter(request => request.url === '/api/history/gpu/prepare')}; +} + +test('saved Resume refreshes readiness before preparation, ignores double clicks, and displays real right-panel progress', async t => { + const env = setup(t), status = deferred(); + env.network.nextStatus = status; + assert.equal(env.resumeGPU.hidden, false); + assert.equal(env.get('st-job-time').hidden, true); + assert.equal(env.get('st-job-progress').dataset.resumable, 'true'); + const firstClick = env.resumeGPU.click(); + const secondClick = env.resumeGPU.click(); + assert.equal(env.resumeGPU.disabled, true); + assert.equal(env.resumeGPU.textContent, 'Resuming…'); + assert.deepEqual(env.requests.map(request => request.url), ['/api/history/status']); + assert.equal(env.prepareRequests().length, 0); + status.resolve(env.response(env.network.deployment)); + await firstClick; await secondClick; await tick(); + assert.equal(env.prepareRequests().length, 1); + assert.equal(env.prepareRequests()[0].headers['X-CSRF-Token'], 'csrf-test'); + assert.equal(JSON.parse(env.prepareRequests()[0].body).operation_id, operationId); + assert.deepEqual(env.requests.slice(0, 2).map(request => request.url), ['/api/history/status', '/api/history/gpu/prepare']); + assert.equal(env.get('st-job-detail').textContent, 'Loading model weights · 40%'); + const meter = env.get('st-job-progress').querySelector('progress'); + assert.equal(meter.hidden, false); assert.equal(meter.value, 40); + assert.equal(env.resumeGPU.hidden, true); + assert.equal(env.get('st-job-progress').dataset.resumable, 'false'); + assert.match(env.get('st-job-time').textContent, /first generation after each Space startup/); + assert.equal(env.get('st-job-time').hidden, false); + assert.equal(env.get('st-native-startup'), null); + assert.equal(env.get('st-gate-note').hidden, true); + assert.equal(env.submissions.length, 0); +}); + +test('status polling after a runtime change repeats CPU preparation once and advances that same requested task', async t => { + const env = setup(t); + await env.resumeGPU.click(); await tick(); + assert.equal(env.prepareRequests().length, 1); + env.network.deployment = deployment('runtime-b', 0, 4); + await env.refreshStatus(); await tick(); + assert.equal(env.prepareRequests().length, 2); + assert.equal(env.get('st-job-detail').textContent, 'Loading model weights · 0%'); + assert.equal(env.resumeGPU.hidden, true); + env.network.deployment = deployment('runtime-b', 3, 4); + await env.refreshStatus(); await tick(); + assert.equal(env.prepareRequests().length, 2); + assert.equal(env.get('st-job-detail').textContent, 'Loading model weights · 75%'); + assert.equal(env.get('st-job-progress').querySelector('progress').value, 75); + env.network.deployment.models = {stage: 'ready_released'}; + await env.refreshStatus(); await tick(); + assert.equal(env.submissions.length, 1); + assert.equal(env.submissions[0].data.operation_id, operationId); + assert.equal(env.get('st-job-detail').textContent, 'Starting generation… Keep this page open.'); + assert.equal(env.get('st-job-time').hidden, true); + assert.equal(env.get('st-job-progress').querySelector('progress').value, undefined); + await env.refreshStatus(); await tick(); + assert.equal(env.prepareRequests().length, 2); + assert.equal(env.submissions.length, 1); +}); + +test('a failed status refresh keeps Resume available and does not prepare until a later successful click', async t => { + const env = setup(t); + env.network.statusFailure = true; + await env.resumeGPU.click(); await tick(); + assert.equal(env.requests.filter(request => request.url === '/api/history/status').length, 2); + assert.equal(env.prepareRequests().length, 0); + assert.equal(env.resumeGPU.hidden, false); + assert.equal(env.resumeGPU.disabled, false); + assert.equal(env.resumeGPU.textContent, 'Resume task'); + assert.equal(env.get('st-job-time').hidden, true); + assert.match(env.get('st-native-notice').textContent, /try Resume task again/); + env.network.statusFailure = false; + await env.resumeGPU.click(); await tick(); + assert.equal(env.prepareRequests().length, 1); + assert.equal(env.resumeGPU.hidden, true); + assert.equal(env.get('st-job-detail').textContent, 'Loading model weights · 40%'); +}); + +function geometryReady(env, fields = {}) { + const job = {...savedJob, workflow: 'geometry_texture', status: 'complete', asset_id: 'white-asset', + texture_status: 'not_requested', gpu_operation_id: null, gpu_operation_kind: null, ...fields}; + env.S.jobs.set(job.id, job); + return job; +} + +test('persisted startup wording is normalized in geometry and texture details without replacing actual stage labels', t => { + const env = setup(t), oldLabel = 'Running the latest single-view pipeline'; + for (const workflow of [undefined, 'geometry_texture']) { + const job = {...savedJob, workflow, status: 'running', gpu_operation_id: null, progress: {label: oldLabel, elapsed_seconds: 12}}; + env.S.jobs.set(job.id, job); env.renderSelected(); + assert.equal(env.get('st-job-detail').textContent, 'Starting generation…'); + assert.equal(env.get('st-job-time').textContent, '12s elapsed'); + job.progress.label = 'Generating geometry'; env.renderSelected(); + assert.equal(env.get('st-job-detail').textContent, 'Generating geometry'); + } + const job = geometryReady(env, {texture_status: 'running', texture_progress: {label: oldLabel}}); + env.S.previewLoaded = env.previewKey(job); env.renderSelected(); + assert.equal(env.get('st-texture-detail').textContent, 'Starting generation…'); + job.texture_progress.label = 'Generating materials'; env.renderSelected(); + assert.equal(env.get('st-texture-detail').textContent, 'Generating materials'); + assert.equal(env.requests.length, 0); +}); + +test('selecting white geometry never starts textures; an explicit click replaces the white display with central texture preparation', async t => { + const env = setup(t), job = geometryReady(env); + env.selectJob(job.id); await tick(); + assert.equal(env.get('st-ovoxel').innerHTML, '
White preview
'); + assert.equal(env.get('st-ovoxel').hidden, false); + assert.equal(env.get('st-texture').textContent, 'Generate textures'); + assert.equal(env.get('st-texture-note').hidden, false); + assert.equal(env.get('st-extract').hidden, true); + assert.equal(env.get('st-extract').disabled, true); + assert.equal(env.prepareRequests().length, 0); + assert.equal(env.requests.filter(request => request.url.endsWith('/texture')).length, 0); + env.network.textureResponse = {...job, texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'b'.repeat(32)}; + const click = env.get('st-texture').click(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + assert.equal(env.get('st-ovoxel').hidden, true); + assert.equal(env.get('st-native-empty').hidden, true); + assert.equal(env.get('st-texture-detail').textContent, 'Connecting to your task…'); + assert.equal(env.get('st-texture-progress').querySelector('progress').hidden, false); + const duplicate = env.get('st-texture').click(); + await click; await duplicate; await tick(); + assert.equal(env.requests.filter(request => request.url.endsWith('/texture')).length, 1); + assert.equal(env.prepareRequests().length, 1); + assert.equal(env.get('st-texture-detail').textContent, 'Loading model weights · 40%'); + assert.equal(env.get('st-texture-progress').querySelector('progress').value, 40); + assert.equal(env.get('st-ovoxel').innerHTML, '
White preview
'); + assert.equal(env.get('st-ovoxel').hidden, true, 'white geometry is retained only for a failure fallback'); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + assert.equal(env.get('st-job-progress').hidden, true); + assert.equal(env.get('st-texture').hidden, true); + assert.equal(env.submissions.length, 0); +}); + +test('texture completion refreshes the asset preview and enables export without changing the selected job', async t => { + const env = setup(t), white = geometryReady(env); + env.selectJob(white.id); await tick(); + const oldKey = env.previewKey(white); + const textured = {...white, asset_id: 'textured-asset', texture_status: 'complete'}; + env.network.previewHTML = '
Textured preview
'; + env.S.jobs.set(white.id, textured); env.renderSelected(); await tick(); + assert.notEqual(env.previewKey(textured), oldKey); + assert.equal(env.S.selected.id, white.id); + assert.equal(env.get('st-ovoxel').innerHTML, '
Textured preview
'); + assert.equal(env.requests.filter(request => request.url.includes('/preview?')).length, 2); + assert.equal(env.get('st-texture-actions').hidden, true); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + assert.equal(env.get('st-ovoxel').hidden, false); + assert.equal(env.get('st-extract').hidden, false); + assert.equal(env.get('st-extract').disabled, false); + assert.equal(env.historyStatus(textured).label, 'Complete'); + assert.equal(env.prepareRequests().length, 0); +}); + +test('saved queued textures resume the existing operation only after a click and failed textures start a new retry', async t => { + const env = setup(t), queued = geometryReady(env, {texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'b'.repeat(32)}); + env.selectJob(queued.id); await tick(); + assert.equal(env.get('st-texture').textContent, 'Resume textures'); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + assert.equal(env.get('st-ovoxel').hidden, true); + assert.equal(env.requests.filter(request => request.url.includes('/preview?')).length, 0); + assert.equal(env.prepareRequests().length, 0); + await env.get('st-texture').click(); await tick(); + assert.equal(env.prepareRequests().length, 1); + assert.equal(env.requests.filter(request => request.url.endsWith('/texture')).length, 0); + const failed = {...queued, texture_status: 'failed', texture_error: 'Texture generation failed.', gpu_operation_id: null}; + env.S.jobs.set(queued.id, failed); env.renderSelected(); await tick(); + assert.equal(env.get('st-texture').textContent, 'Retry textures'); + assert.equal(env.get('st-ovoxel').hidden, false); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + env.network.textureResponse = {...queued, gpu_operation_id: 'c'.repeat(32)}; + await env.get('st-texture').click(); await tick(); + assert.equal(env.requests.filter(request => request.url.endsWith('/texture')).length, 1); + assert.equal(env.prepareRequests().length, 2); + assert.equal(JSON.parse(env.prepareRequests()[1].body).operation_id, 'c'.repeat(32)); +}); + +test('texture POST retries retain their idempotency key while running tasks and legacy jobs need no texture action', async t => { + const env = setup(t), job = geometryReady(env); + env.selectJob(job.id); await tick(); + env.network.textureFailure = true; env.network.textureResponse = {detail: 'Temporarily unavailable'}; + await env.get('st-texture').click(); await tick(); + env.network.textureFailure = false; + env.network.textureResponse = {...job, texture_status: 'running', texture_progress: {label: 'Generating materials', elapsed_seconds: 10}}; + await env.get('st-texture').click(); await tick(); + const posts = env.requests.filter(request => request.url.endsWith('/texture')); + assert.equal(posts.length, 2); + assert.equal(JSON.parse(posts[0].body).idempotency_key, JSON.parse(posts[1].body).idempotency_key); + assert.equal(env.get('st-texture').hidden, true); + assert.equal(env.get('st-texture-title').textContent, 'Generating textures'); + assert.equal(env.get('st-texture-progress').querySelector('progress').value, undefined); + assert.equal(env.historyStatus(env.S.jobs.get(job.id)).label, 'Generating textures'); + const legacy = {...job, workflow: undefined}; + env.S.jobs.set(job.id, legacy); env.renderSelected(); + assert.equal(env.get('st-texture-actions').hidden, true); + assert.equal(env.get('st-extract').hidden, false); + assert.equal(env.historyStatus(legacy).label, 'Generation Complete'); + env.S.examples.set('example', {image_url: '/example.png', preview_html: '
Example preview
'}); + env.selectExample('example'); + assert.equal(env.get('st-model').innerHTML, '
Example preview
'); + assert.equal(env.get('st-texture-actions').hidden, true); +}); + +test('one saved geometry keeps one history card through texture failure, retry and success without submitting geometry again', async t => { + const env = setup(t), white = geometryReady(env, {upload_id: 'saved-upload', seed: 42}); + env.selectJob(white.id); await tick(); + const card = env.get('st-history-list').children[0]; + const sameCase = () => { + assert.equal(env.S.jobs.size, 1); + assert.equal(env.S.selected.id, white.id); + assert.equal(env.get('st-history-list').children.length, 1); + assert.equal(env.get('st-history-list').children[0], card); + assert.equal(card.dataset.jobId, white.id); + assert.equal(env.requests.filter(request => request.url === '/api/history/submit').length, 0); + }; + assert.equal(env.get('st-generate').disabled, true); + assert.equal(env.get('st-generate').textContent, 'Geometry ready'); + assert.equal(env.get('st-input-panel').inert, false); + await env.get('st-generate').click(); + env.get('st-generate').disabled = false; // The handler must guard programmatic bypasses too. + await env.submit(); env.updateGate(); sameCase(); + env.network.textureResponse = {...white, texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'b'.repeat(32)}; + await env.get('st-texture').click(); await tick(); sameCase(); + const failed = {...white, texture_status: 'failed', texture_error: 'Texture generation failed.'}; + env.network.jobsResponse = {items: [failed], revision: 2}; + await env.refreshJobs(); await tick(); sameCase(); + assert.equal(env.get('st-ovoxel').innerHTML, '
White preview
'); + assert.equal(env.get('st-ovoxel').hidden, false); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + assert.equal(env.get('st-texture').textContent, 'Retry textures'); + assert.equal(env.get('st-generate').disabled, true); + env.network.textureResponse = {...white, texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'c'.repeat(32)}; + await env.get('st-texture').click(); await tick(); sameCase(); + const textured = {...white, asset_id: 'textured-asset', texture_status: 'complete'}; + env.network.previewHTML = '
Textured preview
'; + env.network.jobsResponse = {items: [textured], revision: 3}; + await env.refreshJobs(); await tick(); sameCase(); + assert.equal(env.get('st-ovoxel').innerHTML, '
Textured preview
'); + assert.equal(env.get('st-generate').textContent, 'Geometry generated'); + assert.equal(env.get('st-generate').disabled, true); + assert.equal(env.get('st-extract').disabled, false); + assert.equal(env.requests.filter(request => request.url.endsWith('/texture')).length, 2); +}); + +test('changing the selected geometry seed or upload allows a legitimate new geometry submission', async t => { + for (const change of ['seed', 'upload']) { + const env = setup(t), white = geometryReady(env, {upload_id: 'saved-upload', seed: 42}); + env.selectJob(white.id); await tick(); + assert.equal(env.get('st-generate').disabled, true); + if (change === 'seed') { + env.get('st-seed').value = '43'; + for (const handler of env.get('st-seed').handlers.get('input')) handler(); + } else { + env.S.source = {upload_id: 'replacement-upload'}; env.updateGate(); + } + assert.equal(env.get('st-generate').disabled, false); + assert.equal(env.get('st-generate').textContent, 'Generate geometry'); + const submission = deferred(); env.network.nextSubmission = submission; + const pending = env.submit(); + const posts = env.requests.filter(request => request.url === '/api/history/submit'); + assert.equal(posts.length, 1); + const body = JSON.parse(posts[0].body); + assert.equal(body.seed, change === 'seed' ? 43 : 42); + assert.equal(body.upload_id, change === 'upload' ? 'replacement-upload' : 'saved-upload'); + submission.resolve(env.response({...savedJob, id: 'new-job', seed: body.seed, upload_id: body.upload_id, + gpu_operation_id: 'd'.repeat(32)})); + await pending; await tick(); + assert.equal(env.S.jobs.size, 2, 'independent geometry runs must remain separate cases'); + assert.equal(env.S.selected.id, 'new-job'); + } +}); + +test('a mismatched texture response cannot create another card or dispatch its operation', async t => { + const env = setup(t), white = geometryReady(env, {upload_id: 'saved-upload', seed: 42}); + env.selectJob(white.id); await tick(); + env.network.textureResponse = {...white, id: 'wrong-job', texture_status: 'queued', + gpu_operation_kind: 'texture', gpu_operation_id: 'e'.repeat(32)}; + await env.get('st-texture').click(); await tick(); + assert.equal(env.S.jobs.size, 1); + assert.equal(env.get('st-history-list').children.length, 1); + assert.equal(env.S.selected.id, white.id); + assert.equal(env.prepareRequests().length, 0); + assert.equal(env.submissions.length, 0); + assert.equal(env.get('st-ovoxel').innerHTML, '
White preview
'); + assert.match(env.get('st-texture-detail').textContent, /did not match this saved model/); + assert.equal(env.get('st-texture').textContent, 'Retry textures'); +}); + +test('an outdated white preview response is discarded before it changes the visible preview or input', async t => { + const env = setup(t), white = geometryReady(env, {upload_id: 'saved-upload', seed: 42}), preview = deferred(); + const writes = [], panel = env.get('st-ovoxel'); + let html = ''; + Object.defineProperty(panel, 'innerHTML', {get: () => html, set(value) { html = value; writes.push(value); }}); + env.network.nextPreview = preview; + env.selectJob(white.id); await tick(); + const textured = {...white, asset_id: 'textured-asset', texture_status: 'complete'}; + env.S.jobs.set(white.id, textured); + env.network.previewHTML = '
Textured preview
'; + preview.resolve(env.response({view: 'render', html: '
Stale white preview
', + upload_id: 'stale-upload', display_url: '/api/history/jobs/saved-job/files/stale.png'})); + await tick(); await tick(); + assert.deepEqual(writes, ['
Textured preview
']); + assert.equal(env.S.source.upload_id, 'saved-upload'); + assert.equal(env.S.previewLoaded, env.previewKey(textured)); +}); + +test('texture POST and status failures restore cached white geometry with retry, including a cold workspace', async t => { + for (const failure of ['post', 'status', 'cold']) { + const env = setup(t), white = geometryReady(env); + env.selectJob(white.id); await tick(); + if (failure === 'post') { + env.network.textureFailure = true; env.network.textureResponse = {detail: 'Texture request could not connect.'}; + } else if (failure === 'status') env.network.statusFailure = true; + else env.network.deployment = {...env.network.deployment, workspace: {stage: 'starting'}}; + const click = env.get('st-texture').click(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + assert.equal(env.get('st-ovoxel').hidden, true); + await click; await tick(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + assert.equal(env.get('st-ovoxel').hidden, false); + assert.equal(env.get('st-ovoxel').innerHTML, '
White preview
'); + assert.equal(env.get('st-texture').hidden, false); + assert.equal(env.get('st-texture').textContent, 'Retry textures'); + assert.equal(env.get('st-texture').disabled, failure === 'cold'); + assert.equal(env.get('st-texture-title').textContent, 'Texture failed'); + assert.equal(env.get('st-texture-progress').querySelector('progress').hidden, true); + assert.equal(env.requests.filter(request => request.url.includes('/preview?')).length, 1); + } +}); + +test('a dispatcher failure restores white and an actionable retry even while the saved texture is still queued', async t => { + const env = setup(t), white = geometryReady(env); + env.selectJob(white.id); await tick(); + const queued = {...white, texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'b'.repeat(32)}; + env.network.textureResponse = queued; env.network.jobsResponse = {items: [queued], revision: 2}; + await env.get('st-texture').click(); await tick(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + env.gpuDispatcher.advance(() => 'Hugging Face could not verify your personal ZeroGPU quota. Sign in on Hugging Face, then refresh this Space.'); + await tick(); await tick(); + assert.equal(env.S.jobs.get(white.id).texture_status, 'queued'); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + assert.equal(env.get('st-ovoxel').hidden, false); + assert.equal(env.get('st-texture').textContent, 'Retry textures'); + assert.equal(env.get('st-texture').hidden, false); + assert.equal(env.get('st-texture').disabled, false); + assert.equal(env.get('st-texture-progress').querySelector('progress').hidden, true); + assert.match(env.get('st-texture-detail').textContent, /could not verify/); + const retry = env.get('st-texture').click(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + assert.equal(env.get('st-texture-title').textContent, 'Preparing textures'); + assert.equal(env.get('st-ovoxel').hidden, true); + await retry; await tick(); + assert.equal(env.requests.filter(request => request.url.endsWith('/texture')).length, 1); + assert.equal(env.prepareRequests().length, 2); +}); + +test('texture activity belongs to its selected case and cannot hide another history preview after a delayed connection', async t => { + const env = setup(t), white = geometryReady(env), status = deferred(); + env.selectJob(white.id); await tick(); + const other = {...white, id: 'another-job', asset_id: 'another-white'}; + env.S.jobs.set(other.id, other); env.network.previewByJob = {[other.id]: '
Other white preview
'}; + env.network.nextStatus = status; + env.network.textureResponse = {...white, texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'b'.repeat(32)}; + const click = env.get('st-texture').click(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + env.selectJob(other.id); await tick(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + assert.equal(env.get('st-ovoxel').hidden, false); + assert.equal(env.get('st-ovoxel').innerHTML, '
Other white preview
'); + status.resolve(env.response(env.network.deployment)); await click; await tick(); + assert.equal(env.S.selected.id, other.id); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + assert.equal(env.get('st-ovoxel').hidden, false); + const running = {...env.network.textureResponse, texture_status: 'running', texture_progress: {label: 'Generating materials'}}; + env.S.jobs.set(white.id, running); env.selectJob(white.id); await tick(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + assert.equal(env.get('st-ovoxel').hidden, true); + assert.equal(env.get('st-texture-title').textContent, 'Generating textures'); + assert.equal(env.requests.filter(request => request.url.includes('/jobs/saved-job/preview?')).length, 1, + 'returning to active textures must not request white geometry again'); +}); + +test('a pending white preview may be cached during textures but never becomes visible until a failure', async t => { + const env = setup(t), white = geometryReady(env), preview = deferred(); + env.network.nextPreview = preview; env.selectJob(white.id); await tick(); + env.network.textureResponse = {...white, texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'b'.repeat(32)}; + await env.get('st-texture').click(); await tick(); + const panel = env.get('st-ovoxel'), visibility = []; + let hidden = panel.hidden; + Object.defineProperty(panel, 'hidden', {get: () => hidden, set(value) { hidden = value; visibility.push(value); }}); + preview.resolve(env.response({view: 'render', html: '
Late white preview
'})); + await tick(); await tick(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'true'); + assert.equal(panel.hidden, true); + assert.equal(panel.innerHTML, '
Late white preview
'); + assert.equal(visibility.includes(false), false, 'even an intermediate response must not flash white geometry'); + assert.equal(env.get('st-native-empty').hidden, true); + assert.equal(env.S.previewLoaded, env.previewKey(white)); + env.S.jobs.set(white.id, {...white, texture_status: 'failed', texture_error: 'Texture generation failed.'}); + env.renderSelected(); + assert.equal(env.get('st-output-panel').dataset.textureProgress, 'false'); + assert.equal(panel.hidden, false); + assert.equal(env.get('st-texture').textContent, 'Retry textures'); + assert.equal(env.requests.filter(request => request.url.includes('/preview?')).length, 1); +}); + +function assertInputLocked(env, locked) { + assert.equal(env.get('st-input-panel').inert, locked, 'the complete left panel follows generation state'); + for (const id of ['st-file-input', 'st-pick-image', 'st-change-image', 'st-clear-image', 'st-seed', 'st-generate']) + assert.equal(env.get(id).disabled, locked, id); + assert.equal(env.get('st-history-list').inert, false, 'history remains usable'); + assert.equal(env.get('st-output-panel').inert, false, 'right-side progress and resume remain usable'); +} + +test('running geometry prevents input changes, dropped files and repeat submissions until a terminal state', async t => { + const env = setup(t), job = {...savedJob, status: 'running', gpu_operation_id: null}; + env.S.jobs.set(job.id, job); env.get('st-seed').value = '42'; env.renderSelected(); + assertInputLocked(env, true); + const original = env.S.source, selected = env.S.selected, before = env.requests.length; + let pickerOpened = 0; + env.get('st-file-input').addEventListener('click', () => { pickerOpened++; }); + await env.get('st-pick-image').click(); await env.get('st-change-image').click(); + await env.get('st-clear-image').click(); await env.get('st-generate').click(); + for (const handler of env.get('st-upload').handlers.get('drop')) + handler({preventDefault() {}, dataTransfer: {files: [{name: 'replacement.png', size: 10}]}}); + env.S.examples.set('replacement', {image_url: '/replacement.png'}); + env.selectExample('replacement'); await tick(); + let fileNavigationPrevented = false; + for (const handler of env.document.handlers.get('drop') || []) + handler({preventDefault() { fileNavigationPrevented = true; }, dataTransfer: {types: ['Files']}}); + assert.equal(fileNavigationPrevented, true, 'dropping over an inert panel cannot navigate the browser to the file'); + assert.equal(pickerOpened, 0); + assert.equal(env.S.source, original); + assert.equal(env.S.selected, selected); + assert.equal(env.requests.length, before, 'no upload or generation may bypass the disabled panel'); + for (const status of ['failed', 'interrupted']) { + job.status = status; env.renderSelected(); assertInputLocked(env, false); + job.status = 'running'; env.renderSelected(); assertInputLocked(env, true); + } + // A fresh page can still restore the current reference from history while + // keeping its editor locked; browsing must not leave a blank upload panel. + env.S.source = null; env.selectJob(job.id); + assert.match(env.S.source.input_url, /\/files\/input\.png$/); + assertInputLocked(env, true); + const white = geometryReady(env); + env.S.previewLoaded = env.previewKey(white); env.renderSelected(); + assertInputLocked(env, false); + assert.equal(env.get('st-texture').disabled, false, 'geometry completion exposes the next explicit stage'); + let idleDropPrevented = false; + for (const handler of env.document.handlers.get('drop') || []) + handler({preventDefault() { idleDropPrevented = true; }, dataTransfer: {types: ['Files']}}); + assert.equal(idleDropPrevented, false, 'idle upload drop handling remains available'); +}); + +test('Resume locks immediately through preparation and releases after a failed GPU connection', async t => { + const env = setup(t), status = deferred(); + assertInputLocked(env, false); // A saved, unrequested task must not lock the workspace forever. + env.network.nextStatus = status; + const click = env.resumeGPU.click(); + assertInputLocked(env, true); + status.resolve(env.response(env.network.deployment)); + await click; await tick(); + assert.equal(env.prepareRequests().length, 1); + assertInputLocked(env, true); + env.gpuDispatcher.advance(() => 'Model preparation failed'); + await tick(); + assertInputLocked(env, false); + assert.equal(env.resumeGPU.hidden, false); + assert.equal(env.resumeGPU.disabled, false); +}); + +test('texture generation locks the input across history selection and polling unlocks when it completes', async t => { + const env = setup(t), white = geometryReady(env), status = deferred(); + env.S.previewLoaded = env.previewKey(white); env.renderSelected(); assertInputLocked(env, false); + env.network.nextStatus = status; + env.network.textureResponse = {...white, texture_status: 'queued', gpu_operation_kind: 'texture', gpu_operation_id: 'b'.repeat(32)}; + const click = env.get('st-texture').click(); assertInputLocked(env, true); + status.resolve(env.response(env.network.deployment)); await click; await tick(); + assertInputLocked(env, true); + const active = {...env.network.textureResponse, texture_status: 'running'}; + env.S.jobs.set(active.id, active); + const older = {...white, id: 'older', asset_id: 'older-white'}; + env.S.jobs.set(older.id, older); + const source = env.S.source; + env.selectJob(older.id); await tick(); + assert.equal(env.S.selected.id, older.id); + assert.equal(env.S.source, source, 'history browsing must not replace the locked reference image'); + assertInputLocked(env, true); + env.network.jobsResponse = {items: [{...active, texture_status: 'complete', gpu_operation_id: null}], revision: 2}; + await env.refreshJobs(); await tick(); + assertInputLocked(env, false); +}); + +test('the first submission locks inputs before its response and rejects a second click', async t => { + const env = setup(t), submission = deferred(); + env.get('st-seed').value = '42'; env.network.nextSubmission = submission; + const first = env.submit(); + assertInputLocked(env, true); + await env.submit(); + assert.equal(env.requests.filter(request => request.url === '/api/history/submit').length, 1); + const queued = {...savedJob, id: 'new-job', gpu_operation_id: 'c'.repeat(32)}; + submission.resolve(env.response(queued)); await first; await tick(); + assertInputLocked(env, true); + env.network.jobsResponse = {items: [{...queued, status: 'failed', gpu_operation_id: null}], revision: 2}; + await env.refreshJobs(); await tick(); + assertInputLocked(env, false); +}); + + +function uploadResult(id = 'upload-a', version) { + return {upload_id: id, foreground: {status: 'complete'}, + foreground_url: '/api/history/uploads/' + id + '/files/foreground.png', + ...(version ? {files: {'foreground.png': {version}}} : {})}; +} +function uploading(env, id = 'upload-a') { + env.clearUploadPreview(); + const state = env.S.uploadState = {active: true, inputVersion: ++env.S.inputVersion, + identityVersion: env.S.identityVersion, controller: new AbortController(), uploadId: id, + timer: null, resultURL: null, failures: 0, phase: 'running'}; + env.S.source = {upload_id: id, foreground_required: true, foreground_ready: false}; + env.S.uploadBusy = true; + return state; +} +const cutoutRequests = env => env.requests.filter(request => request.url.endsWith('/files/foreground.png')); + +test('generation status polling never reloads an already applied foreground or flashes its completed label', async t => { + const env = setup(t), state = uploading(env), result = uploadResult(); + await env.applyUploadResult(result, state); + const source = env.S.source, src = env.get('st-input-image').src; + const labels = [...env.get('st-upload-label').textHistory]; + env.S.submitKey = 'submission-in-flight'; + env.S.jobs.set(savedJob.id, {...savedJob, status: 'running', gpu_operation_id: null}); + for (let i = 0; i < 3; i++) { await env.refreshStatus(); await env.applyUploadResult({...result}, state); } + await tick(); + assert.equal(cutoutRequests(env).length, 1); + assert.equal(env.images.decoded, 1); + assert.equal(env.S.source, source); assert.equal(env.get('st-input-image').src, src); + assert.equal(env.S.submitKey, 'submission-in-flight', 'repeated results must not reset submission identity'); + assert.equal(env.S.source.foreground_ready, true); assert.equal(env.S.uploadBusy, false); + assert.deepEqual(env.get('st-upload-label').textHistory, labels); + assert.equal(env.get('st-upload-label').textContent, 'Background removed.'); + assert.equal(env.requests.some(request => /\/uploads\/upload-a$/.test(request.url)), false); + assert.equal(env.requests.some(request => request.method === 'POST'), false); + assert.equal(env.submissions.length, 0); +}); + +test('overlapping upload polling and completion callbacks share one image fetch and decode', async t => { + const env = setup(t), state = uploading(env), result = uploadResult(), image = deferred(); + env.network.uploadResponse = result; env.network.nextImage = image; + const first = env.applyUploadResult(result, state); + const refreshA = env.refreshUpload(state), refreshB = env.refreshUpload(state); + const callback = env.applyUploadResult({...result}, state); + await tick(); + assert.equal(cutoutRequests(env).length, 1); + assert.equal(env.requests.filter(request => /\/uploads\/upload-a$/.test(request.url)).length, 1); + image.resolve(new Response(new Blob(['cutout'], {type: 'image/png'}))); + await Promise.all([first, refreshA, refreshB, callback]); + assert.equal(env.images.decoded, 1); assert.equal(env.images.created.length, 1); + assert.equal(env.get('st-upload-label').textHistory.filter(text => text === 'Loading image…').length, 1); + assert.equal(env.get('st-upload-label').textHistory.filter(text => text === 'Background removed.').length, 1); + assert.equal(state.timer, null); +}); + +test('a disconnected cutout download can recover, then explicit foreground retry can apply its new result', async t => { + const env = setup(t), state = uploading(env), result = uploadResult(); + env.network.uploadResponse = result; env.network.imageFailure = true; + await env.refreshUpload(state); + assert.equal(state.appliedResultKey, undefined); + assert.equal(env.get('st-upload-label').textContent, 'Preparing image…'); + env.network.imageFailure = false; + await env.refreshUpload(state); + assert.equal(env.S.source.foreground_ready, true); assert.equal(env.images.decoded, 1); + const retried = uploading(env, 'upload-retry'); + await env.applyUploadResult({upload_id: 'upload-retry', foreground: {status: 'failed', error: 'Please retry.'}}, retried); + assert.equal(env.get('st-upload-retry').hidden, false); + env.network.retryUploadResponse = uploadResult('upload-retry'); + await env.retryUploadPreview(); + assert.equal(env.images.decoded, 2); + assert.equal(env.S.source.upload_id, 'upload-retry'); assert.equal(env.S.source.foreground_ready, true); + assert.equal(env.requests.filter(request => request.method === 'POST' && request.url.endsWith('/foreground')).length, 1); +}); + +test('a different result version and a newly selected upload each apply once', async t => { + const env = setup(t), state = uploading(env); + await env.applyUploadResult(uploadResult('upload-a', 'version-1'), state); + const previous = env.get('st-input-image').src; + await env.applyUploadResult(uploadResult('upload-a', 'version-2'), state); + await env.applyUploadResult(uploadResult('upload-a', 'version-2'), state); + assert.notEqual(env.get('st-input-image').src, previous); + assert.equal(env.images.decoded, 2); assert.ok(env.images.revoked.includes(previous)); + const next = uploading(env, 'upload-b'); + await env.applyUploadResult(uploadResult('upload-b'), next); + await env.applyUploadResult(uploadResult('upload-b'), next); + assert.equal(env.images.decoded, 3); assert.equal(env.S.source.upload_id, 'upload-b'); +}); + +test('a foreground decode from a previous account or input cannot replace the current image', {timeout: 2000}, async t => { + for (const change of ['account', 'input']) { + const env = setup(t), old = uploading(env), decode = deferred(), started = deferred(); + env.network.nextDecode = decode; env.network.onDecode = started.resolve; + const pending = env.applyUploadResult(uploadResult(), old); + await started.promise; + const staleURL = old.resultURL; + if (change === 'account') env.clearPrivate(); + const current = uploading(env, 'upload-new'); + await env.applyUploadResult(uploadResult('upload-new'), current); + const source = env.S.source, src = env.get('st-input-image').src; + decode.resolve(); await pending; + assert.equal(env.S.source, source); assert.equal(env.get('st-input-image').src, src); + assert.equal(env.S.source.upload_id, 'upload-new'); + assert.ok(env.images.revoked.includes(staleURL)); + assert.equal(old.appliedResultKey, undefined); + assert.equal(env.get('st-upload-label').textContent, 'Background removed.'); + } +}); + + +test('late queued, failed or unavailable upload status cannot replace an already completed foreground', async t => { + for (const late of ['queued', 'failed', 'unavailable']) { + const env = setup(t), state = uploading(env), pending = deferred(); + env.network.nextUploadStatus = pending; + const refresh = env.refreshUpload(state); + await env.applyUploadResult(uploadResult(), state); + const source = env.S.source, labels = [...env.get('st-upload-label').textHistory]; + pending.resolve(env.response({upload_id: 'upload-a', foreground: {status: late}}, late === 'unavailable' ? 400 : 200)); + await refresh; + assert.equal(env.S.source, source); assert.equal(env.S.source.foreground_ready, true); + assert.equal(env.S.uploadBusy, false); assert.equal(env.get('st-upload-retry').hidden, true); + assert.deepEqual(env.get('st-upload-label').textHistory, labels); + assert.equal(cutoutRequests(env).length, 1); assert.equal(state.timer, null); + } +}); + + +test('a status read failure during cutout decode cannot revoke or replace the image being applied', {timeout: 2000}, async t => { + const env = setup(t), state = uploading(env), decode = deferred(), started = deferred(), status = deferred(); + env.network.nextDecode = decode; env.network.onDecode = started.resolve; + const apply = env.applyUploadResult(uploadResult(), state); + await started.promise; + const resultURL = state.resultURL; + env.network.nextUploadStatus = status; + const refresh = env.refreshUpload(state); + status.resolve(env.response({detail: 'Status temporarily unavailable'}, 400)); + await refresh; + assert.equal(state.resultURL, resultURL); + assert.equal(env.images.revoked.includes(resultURL), false); + assert.equal(env.get('st-upload-label').textContent, 'Loading image…'); + decode.resolve(); await apply; + assert.equal(env.S.localURL, resultURL); assert.equal(env.get('st-input-image').src, resultURL); + assert.equal(env.S.source.foreground_ready, true); assert.equal(env.S.uploadBusy, false); + assert.equal(env.get('st-upload-label').textContent, 'Background removed.'); +}); diff --git a/tests/test_native_eigen.py b/tests/test_native_eigen.py new file mode 100644 index 0000000000000000000000000000000000000000..aa1cdd6be80b4bb10ef61752d4f11c28b8b79728 --- /dev/null +++ b/tests/test_native_eigen.py @@ -0,0 +1,83 @@ +"""The pinned Eigen dependency is available without a GitLab submodule fetch.""" +import hashlib +import io +import tarfile +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from hf_space import native_build + + +def test_vendored_archive_matches_pinned_checksum_and_extracts_headers(tmp_path): + archive = native_build.EIGEN_ARCHIVE + assert archive.stat().st_size == 3025643 + assert hashlib.sha256(archive.read_bytes()).hexdigest() == native_build.EIGEN_SHA256 + checkout = tmp_path / 'TRELLIS.2' + checkout.mkdir() + outside = tmp_path / 'keep.txt' + outside.write_text('unrelated workspace file') + target = native_build.prepare_eigen(checkout) + assert target == checkout / 'o-voxel' / 'third_party' / 'eigen' + assert b'Eigen' in (target / 'Eigen' / 'Core').read_bytes() + assert (target / 'Eigen' / 'src' / 'Core' / 'Matrix.h').is_file() + assert 'Mozilla Public License' in (target / 'COPYING.MPL2').read_text() + assert (target / 'COPYING.README').is_file() + assert outside.read_text() == 'unrelated workspace file' + + +def test_bad_checksum_never_creates_or_changes_eigen_tree(tmp_path): + archive = tmp_path / 'invalid.tar.gz' + archive.write_bytes(b'corrupt download') + checkout = tmp_path / 'checkout' + checkout.mkdir() + with pytest.raises(RuntimeError, match='checksum mismatch'): + native_build.prepare_eigen(checkout, archive) + assert list(checkout.iterdir()) == [] + + +@pytest.mark.parametrize('kind', ['traversal', 'symlink', 'hardlink']) +def test_tar_traversal_and_links_are_rejected_before_any_extraction(tmp_path, monkeypatch, kind): + root = f'eigen-{native_build.EIGEN_REVISION}' + archive = tmp_path / 'malicious.tar.gz' + with tarfile.open(archive, 'w:gz') as output: + first = tarfile.TarInfo(root + '/normal.h') + first.size = 6 + output.addfile(first, io.BytesIO(b'normal')) + member = tarfile.TarInfo(root + ('/../../escaped' if kind == 'traversal' else '/link')) + if kind != 'traversal': + member.type = tarfile.SYMTYPE if kind == 'symlink' else tarfile.LNKTYPE + member.linkname = '../../escaped' + output.addfile(member) + monkeypatch.setattr(native_build, 'EIGEN_SHA256', hashlib.sha256(archive.read_bytes()).hexdigest()) + checkout = tmp_path / 'checkout' + checkout.mkdir() + with pytest.raises(RuntimeError, match='Unsafe member'): + native_build.prepare_eigen(checkout, archive) + assert list(checkout.iterdir()) == [] + assert not (tmp_path / 'escaped').exists() + + +def test_only_pinned_trellis_uses_local_eigen_other_repos_keep_submodules(tmp_path, monkeypatch): + extract, run = Mock(), Mock() + monkeypatch.setattr(native_build, 'prepare_eigen', extract) + monkeypatch.setattr(native_build, 'run', run) + native_build.prepare_submodules('TRELLIS.2', native_build.EIGEN_TRELLIS_REVISION, tmp_path, {}) + extract.assert_called_once_with(tmp_path) + run.assert_not_called() + native_build.prepare_submodules('CuMesh', 'other-revision', tmp_path, {'BUILD_TARGET': 'cuda'}) + assert run.call_args.args[0] == ['git', '-C', tmp_path, 'submodule', 'update', '--init', '--recursive', '--depth=1'] + with pytest.raises(RuntimeError, match='Review the vendored Eigen revision'): + native_build.prepare_submodules('TRELLIS.2', 'different-revision', tmp_path, {}) + + +def test_capacity_diagnostics_report_disk_and_cgroup_limit_only(tmp_path, monkeypatch): + (tmp_path / 'memory.max').write_text('2147483648\n') + monkeypatch.setattr(native_build.shutil, 'disk_usage', lambda _: SimpleNamespace(total=1000, free=600)) + assert native_build.startup_resources(tmp_path, tmp_path) == { + 'tmp_disk_total_bytes': 1000, 'tmp_disk_free_bytes': 600, + 'cgroup_memory_limit_bytes': 2147483648, + } + (tmp_path / 'memory.max').write_text('max\n') + assert native_build.startup_resources(tmp_path, tmp_path)['cgroup_memory_limit_bytes'] is None diff --git a/tests/test_observation_filter_export.py b/tests/test_observation_filter_export.py new file mode 100644 index 0000000000000000000000000000000000000000..8141943ef93141e35a518438b25a42c2367ac5f5 --- /dev/null +++ b/tests/test_observation_filter_export.py @@ -0,0 +1,115 @@ +"""Completion's depth admission survives storage and deferred texture refinement.""" +import json +from pathlib import Path +import sys +from types import SimpleNamespace + +import numpy as np +from PIL import Image +import pytest + +from hf_space import ovoxel_results as results, zerogpu_runtime + + +def reference_inputs(folder): + folder.mkdir(parents=True) + (folder/'camera.json').write_text(json.dumps({'camera_angle_x': .5, 'distance': 2})) + Image.new('RGBA', (8, 8), 'white').save(folder/'condition_rgba.png') + for name in ('albedo.png', 'mra.png'): + Image.new('RGB', (1024, 1024), 'white').save(folder/name) + return folder + + +def filter_outputs(asset): + keep = np.ones((1024, 1024), dtype=np.bool_) + keep[100:200, 250:400] = False + np.save(asset.with_suffix('.observation_keep.npy'), keep, allow_pickle=False) + asset.with_suffix('.observation_admission.json').write_text(json.dumps({'enabled': True})) + return keep + + +def test_mask_survives_durable_reference_restore_without_geometry_changes(tmp_path, monkeypatch): + job = tmp_path/'job' + reference = reference_inputs(job/'asset_work'/'input') + asset = job/'asset.npz' + asset.write_bytes(b'immutable-final-geometry-and-material') + (reference/'geometry.pt').write_bytes(b'immutable-original-white-geometry') + expected = filter_outputs(asset) + results._stage_observation_filter(asset, reference) + durable = tmp_path/'durable' + durable.mkdir() + results._persist_reference_inputs(reference, durable) + manifest = json.loads((durable/results.REFERENCE_DIRECTORY/'manifest.json').read_text()) + assert set(results.OBSERVATION_FILES) <= manifest['files'].keys() + monkeypatch.setattr(results, 'JOBS', tmp_path/'restarted-worker') + monkeypatch.setattr(results, '_copy_progress', lambda src, dst, token: dst.write_bytes(src.read_bytes())) + work = tmp_path/'export' + work.mkdir() + restored = results._reference_inputs('a'*32, durable, work, None) + np.testing.assert_array_equal(results._observation_keep(restored), expected) + assert asset.read_bytes() == b'immutable-final-geometry-and-material' + assert (reference/'geometry.pt').read_bytes() == b'immutable-original-white-geometry' + + +def test_legacy_references_do_not_require_a_filter(tmp_path): + reference = reference_inputs(tmp_path/'legacy') + results._stage_observation_filter(tmp_path/'asset.npz', reference) + assert results._observation_keep(reference) is None + assert set(results._validate_reference_inputs(reference)) == set(results.REFERENCE_FILES) + + +@pytest.mark.parametrize('missing', ['observation_keep.npy', 'observation_admission.json']) +def test_partial_filter_cannot_silently_reproject_unfiltered_colors(tmp_path, missing): + reference = reference_inputs(tmp_path/'reference') + asset = tmp_path/'asset.npz' + filter_outputs(asset) + asset.with_suffix('.'+missing).unlink() + with pytest.raises(ValueError, match='did not save'): + results._stage_observation_filter(asset, reference) + + +def test_new_depth_completion_requires_saved_filter(tmp_path): + reference = reference_inputs(tmp_path/'reference') + np.save(reference/'image_depth.npy', np.ones((8, 8), dtype=np.float32)) + with pytest.raises(ValueError, match='did not save'): + results._stage_observation_filter(tmp_path/'asset.npz', reference) + + +def test_manifest_bound_mask_rejects_changes_and_disappearance(tmp_path): + reference = reference_inputs(tmp_path/'reference') + asset = tmp_path/'asset.npz' + filter_outputs(asset) + results._stage_observation_filter(asset, reference) + bindings = results._validate_reference_inputs(reference) + np.save(reference/'observation_keep.npy', np.ones((1024, 1024), dtype=bool)) + with pytest.raises(ValueError, match='integrity'): + results._validate_reference_inputs(reference, bindings) + (reference/'observation_keep.npy').unlink() + with pytest.raises(ValueError, match='incomplete'): + results._validate_reference_inputs(reference, bindings) + + +@pytest.mark.parametrize('mask', [np.ones((8, 8), dtype=bool), np.ones((1024, 1024), dtype=np.uint8)]) +def test_bad_mask_dimensions_and_types_are_rejected(tmp_path, mask): + np.save(tmp_path/'observation_keep.npy', mask) + (tmp_path/'observation_admission.json').write_text('{}') + with pytest.raises(ValueError, match='mask is invalid'): + results._observation_keep(tmp_path) + + +def test_zerogpu_deferred_export_passes_original_admission_mask(tmp_path, monkeypatch): + reference = reference_inputs(tmp_path/'reference') + asset = tmp_path/'asset.npz' + expected = filter_outputs(asset) + results._stage_observation_filter(asset, reference) + calls = [] + monkeypatch.setattr(results, '_export_local', lambda *args, **kwargs: None) + def reproject(output, source, work, *, stage, keep): + assert source == reference + np.testing.assert_array_equal(keep, expected) + calls.append(output) + return output + monkeypatch.setitem(sys.modules, 'supratex.inference.reprojection', SimpleNamespace(run_front_view=reproject)) + output = tmp_path/'asset.glb' + value, _ = zerogpu_runtime.export(str(asset), str(output), str(reference), str(tmp_path/'work'), None) + assert value == str(output) and calls == [output] diff --git a/tests/test_zerogpu_client.cjs b/tests/test_zerogpu_client.cjs new file mode 100644 index 0000000000000000000000000000000000000000..73d0dcdb01c1a3a75c434417a620585acc59e4fe --- /dev/null +++ b/tests/test_zerogpu_client.cjs @@ -0,0 +1,631 @@ +const test = require('node:test'); +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const path = require('node:path'); +const vm = require('node:vm'); +const source = fs.readFileSync(path.join(__dirname, '../hf_space/zerogpu_client.js'), 'utf8'); +const id = 'a'.repeat(32); +const snapshot = {gpu_mode: 'zerogpu', gpu_operation_id: id}; +const deferred = () => { let resolve, reject; const promise = new Promise((r, j) => { resolve = r; reject = j; }); return {promise, resolve, reject}; }; +const tick = () => new Promise(resolve => setImmediate(resolve)); +function setup(connect, overrides = {}, environment = {}) { + const context = vm.createContext({AbortController, Headers, URL, setTimeout, clearTimeout, + location: {origin: 'https://example.hf.space'}, ...environment}); + vm.runInContext(source, context); + const state = {identity: {authenticated: true, csrf_token: 'csrf-1'}, version: 0, stopped: false, ready: true}; + const errors = [], completed = []; + const dispatcher = context.SupraZeroGPU.createDispatcher({ + identity: () => state.identity, version: () => state.version, stopped: () => state.stopped, ready: () => state.ready, + onError: (record, message) => errors.push({record, message}), + onComplete: (record, result) => completed.push({record, result}), connect, + quotaHeaders: async () => ({'x-zerogpu-token': 'personal-quota-token', 'x-zerogpu-uuid': 'session-id'}), ...overrides + }); + return {dispatcher, state, errors, completed, runtime: context.SupraZeroGPU}; +} +function client(events, calls) { + return {submit(endpoint, data, eventData, trigger, allEvents, headers) { + calls.push({endpoint, data, headers, eventData, trigger, allEvents}); + return (async function* () { for (const event of events) yield event; })(); + }}; +} + +test('only authenticated ZeroGPU snapshots can dispatch, never anonymous or dedicated GPU work', () => { + let connections = 0; + const {dispatcher, state} = setup(() => { connections++; }); + dispatcher.observe({gpu_operation_id: id}); + dispatcher.observe({...snapshot, gpu_operation_id: '../../someone-else'}); + state.identity.authenticated = false; + dispatcher.observe(snapshot); + assert.equal(connections, 0); +}); + +test('polling submits each operation once and forwards CSRF without owner credentials', async () => { + const calls = [], connection = deferred(), options = []; + const {dispatcher, completed} = setup(config => { options.push(config); return connection.promise; }); + dispatcher.observe(snapshot, {requested: true}); + dispatcher.observe({items: [snapshot, snapshot]}); + const finished = dispatcher.dispatch(snapshot); + connection.resolve(client([{type: 'data', data: [{status: 'complete'}]}], calls)); + await finished; + dispatcher.observe(snapshot); + assert.equal(options.length, 1); + assert.deepEqual(JSON.parse(JSON.stringify(options[0])), { + events: ['data', 'status'], headers: {'X-CSRF-Token': 'csrf-1'} + }); + assert.equal(calls.length, 1); + assert.equal(calls[0].endpoint, '/run_gpu_operation'); + assert.equal(calls[0].data.operation_id, id); + assert.equal(calls[0].headers['x-zerogpu-token'], 'personal-quota-token'); + assert.equal(calls[0].headers['x-zerogpu-uuid'], 'session-id'); + assert.equal(calls[0].headers['x-gradio-user'], 'app'); + assert.equal(calls[0].allEvents, undefined); + assert.equal(completed.length, 1); +}); + +test('explicit task prepares once, waits for models, then polling advances its first GPU attempt', async () => { + const calls = [], prepared = []; + const {dispatcher, state} = setup(() => client([], calls), {prepare: async record => prepared.push(record.gpu_operation_id)}); + state.ready = false; + await dispatcher.dispatch(snapshot); + dispatcher.observe(snapshot); + assert.deepEqual(prepared, [id]); + assert.equal(calls.length, 0); + assert.equal(dispatcher.pending(), true); + state.ready = true; + dispatcher.advance(); + await dispatcher.dispatch(snapshot); + assert.equal(calls.length, 1); + assert.equal(dispatcher.pending(), false); +}); + +test('opening and polling saved history never prepares models or dispatches old queued work', async () => { + const calls = [], prepared = []; + const {dispatcher} = setup(() => client([], calls), {prepare: async record => prepared.push(record.gpu_operation_id)}); + dispatcher.observe(snapshot); dispatcher.observe({items: [snapshot]}); dispatcher.advance(); + await tick(); + assert.equal(dispatcher.requested(id), false); + assert.equal(prepared.length, 0); assert.equal(calls.length, 0); + await dispatcher.retry(snapshot); + assert.deepEqual(prepared, [id]); assert.equal(calls.length, 1); +}); + +test('the first runtime initializes without repeating preparation and unchanged runtime polls are inert', async () => { + const calls = [], prepared = []; + const {dispatcher, state} = setup(() => client([], calls), {prepare: async record => prepared.push(record.gpu_operation_id)}); + state.ready = false; + await dispatcher.dispatch(snapshot); + dispatcher.syncRuntime('runtime-a'); + dispatcher.syncRuntime('runtime-a'); + dispatcher.observe(snapshot); + await tick(); + assert.deepEqual(prepared, [id]); + assert.equal(calls.length, 0); + assert.equal(dispatcher.pending(), true); +}); + +test('a new runtime repeats CPU preparation only for explicitly requested waiting work', async () => { + const calls = [], prepared = []; + const saved = {...snapshot, gpu_operation_id: 'b'.repeat(32)}; + const {dispatcher, state} = setup(() => client([], calls), {prepare: async record => prepared.push(record.gpu_operation_id)}); + state.ready = false; + dispatcher.syncRuntime('runtime-a'); + dispatcher.observe(saved); + await dispatcher.dispatch(snapshot); + dispatcher.syncRuntime('runtime-b'); + await tick(); + dispatcher.syncRuntime('runtime-b'); + dispatcher.observe({items: [snapshot, saved]}); + await tick(); + assert.deepEqual(prepared, [id, id]); + assert.equal(dispatcher.requested(saved.gpu_operation_id), false); + assert.equal(calls.length, 0); + state.ready = true; + dispatcher.advance(); + await dispatcher.dispatch(snapshot); + assert.deepEqual(calls.map(call => call.data.operation_id), [id]); +}); + +test('an earlier runtime preparation success cannot bypass the new runtime preparation', async () => { + const calls = [], attempts = [], oldPrepare = deferred(), newPrepare = deferred(); + const {dispatcher, state, errors, completed} = setup(() => client([], calls), { + prepare(record, signal) { + attempts.push({record, signal}); + return attempts.length === 1 ? oldPrepare.promise : newPrepare.promise; + } + }); + state.ready = false; + dispatcher.syncRuntime('runtime-a'); + const stale = dispatcher.dispatch(snapshot); + dispatcher.syncRuntime('runtime-b'); + await tick(); + assert.equal(attempts.length, 2); + state.ready = true; + oldPrepare.resolve(); + await stale; await tick(); + assert.equal(calls.length, 0); + assert.equal(errors.length, 0); + assert.equal(completed.length, 0); + assert.equal(dispatcher.pending(), true); + newPrepare.resolve(); + await dispatcher.dispatch(snapshot); + assert.equal(calls.length, 1); + assert.equal(completed.length, 1); +}); + +test('an earlier runtime preparation failure cannot fail or cancel the current attempt', async () => { + const calls = [], oldPrepare = deferred(), newPrepare = deferred(); + let attempts = 0; + const {dispatcher, state, errors, completed} = setup(() => client([], calls), { + prepare: () => ++attempts === 1 ? oldPrepare.promise : newPrepare.promise + }); + state.ready = false; + dispatcher.syncRuntime('runtime-a'); + const stale = dispatcher.dispatch(snapshot); + dispatcher.syncRuntime('runtime-b'); + await tick(); + assert.equal(attempts, 2); + oldPrepare.reject(new Error('Model preparation failed in the previous runtime')); + await stale; await tick(); + assert.equal(errors.length, 0); + assert.equal(dispatcher.failure(id), ''); + assert.equal(dispatcher.pending(), true); + state.ready = true; + newPrepare.resolve(); + await dispatcher.dispatch(snapshot); + assert.equal(calls.length, 1); + assert.equal(completed.length, 1); +}); + +test('runtime changes never replay running, completed, or failed GPU submissions', async () => { + const calls = [], prepared = [], finishRunning = deferred(); + const running = {...snapshot, gpu_operation_id: 'b'.repeat(32)}; + const failed = {...snapshot, gpu_operation_id: 'c'.repeat(32)}; + let connections = 0; + const {dispatcher, state, errors, completed} = setup(() => { + connections++; + if (connections === 2) return {submit(endpoint, data) { + calls.push({endpoint, data}); + return (async function* () { + await finishRunning.promise; + yield {type: 'data', data: [{status: 'complete'}]}; + })(); + }}; + return client(connections === 3 ? [{type: 'status', stage: 'error', message: 'GPU quota exceeded'}] : [], calls); + }, {prepare: async record => prepared.push(record.gpu_operation_id)}); + dispatcher.syncRuntime('runtime-a'); + await dispatcher.dispatch(snapshot); + const runningTask = dispatcher.dispatch(running); + await tick(); + await dispatcher.dispatch(failed); + state.ready = false; + dispatcher.syncRuntime('runtime-b'); + dispatcher.observe({items: [snapshot, running, failed]}); + await tick(); + assert.deepEqual(prepared, [id, running.gpu_operation_id, failed.gpu_operation_id]); + assert.equal(calls.length, 3); + assert.equal(errors.length, 1); + finishRunning.resolve(); + await runningTask; + assert.equal(completed.length, 2); + dispatcher.syncRuntime('runtime-c'); + await tick(); + assert.equal(calls.length, 3); + assert.equal(prepared.length, 3); +}); + +test('a runtime change does not restart a GPU connection already in progress', async () => { + const calls = [], prepared = [], connection = deferred(); + let connections = 0; + const {dispatcher, state, completed} = setup(() => { connections++; return connection.promise; }, { + prepare: async record => prepared.push(record.gpu_operation_id) + }); + dispatcher.syncRuntime('runtime-a'); + const connecting = dispatcher.dispatch(snapshot); + await tick(); + assert.equal(connections, 1); + state.ready = false; + dispatcher.syncRuntime('runtime-b'); + await tick(); + connection.resolve(client([], calls)); + await connecting; + assert.deepEqual(prepared, [id]); + assert.equal(connections, 1); + assert.equal(calls.length, 1); + assert.equal(completed.length, 1); +}); + +test('account reset prevents runtime changes from replaying waiting or unresolved preparation', async () => { + const calls = [], prepared = [], oldPrepare = deferred(); + const waiting = {...snapshot, gpu_operation_id: 'b'.repeat(32)}; + const {dispatcher, state, errors, completed} = setup(() => client([], calls), { + prepare(record) { + prepared.push(record.gpu_operation_id); + return record.gpu_operation_id === id ? oldPrepare.promise : Promise.resolve(); + } + }); + state.ready = false; + dispatcher.syncRuntime('runtime-a'); + const stale = dispatcher.dispatch(snapshot); + await dispatcher.dispatch(waiting); + state.version++; + state.identity = {authenticated: true, csrf_token: 'new-account'}; + dispatcher.reset(); + dispatcher.syncRuntime('runtime-b'); + dispatcher.observe({items: [snapshot, waiting]}); + state.ready = true; + oldPrepare.resolve(); + await stale; + dispatcher.syncRuntime('runtime-c'); + dispatcher.advance(); + await tick(); + assert.deepEqual(prepared, [id, waiting.gpu_operation_id]); + assert.equal(dispatcher.requested(id), false); + assert.equal(dispatcher.requested(waiting.gpu_operation_id), false); + assert.equal(calls.length, 0); + assert.equal(errors.length, 0); + assert.equal(completed.length, 0); +}); + +test('background model failure releases the waiting task for an explicit retry', async () => { + const calls = []; + const {dispatcher, state, errors} = setup(() => client([], calls)); + state.ready = false; + await dispatcher.dispatch(snapshot); + dispatcher.advance('Model preparation failed'); + dispatcher.observe(snapshot); + assert.equal(calls.length, 0); assert.equal(dispatcher.pending(), false); + assert.match(errors[0].message, /Models could not be prepared/); + state.ready = true; + await dispatcher.retry(snapshot); + assert.equal(calls.length, 1); +}); + +test('foreground dispatch waits only for its model and ignores failures in full generation preparation', async () => { + const calls = [], foreground = {...snapshot, gpu_operation_kind: 'foreground'}; + const generation = {...snapshot, gpu_operation_id: 'b'.repeat(32), gpu_operation_kind: 'generation'}; + const deployment = {gpu_mode: 'zerogpu', workspace: {stage: 'ready'}, foreground: {stage: 'loading_models'}, + native_build: {stage: 'compiled'}, models: {stage: 'failed'}}; + const env = setup(() => client([], calls), {ready: record => env.runtime.operationReady(deployment, record)}); + await env.dispatcher.dispatch(foreground); + await env.dispatcher.dispatch(generation); + env.dispatcher.advance(record => env.runtime.preparationError(deployment, record)); + assert.equal(env.errors.length, 1); + assert.equal(env.errors[0].record.gpu_operation_id, generation.gpu_operation_id); + assert.equal(env.dispatcher.failure(id), ''); + assert.equal(env.dispatcher.pending(record => !env.runtime.operationReady(deployment, record)), true); + deployment.foreground.stage = 'ready_released'; + env.dispatcher.advance(record => env.runtime.preparationError(deployment, record)); + await env.dispatcher.dispatch(foreground); + assert.deepEqual(calls.map(call => call.data.operation_id), [id]); + assert.equal(env.dispatcher.pending(), false); +}); + +test('quota errors are sanitized and polling does not repeatedly spend or retry quota', async () => { + const calls = []; + let attempts = 0; + const {dispatcher, errors, completed} = setup(() => { + attempts++; + return client([{type: 'status', stage: 'error', message: 'GPU quota exceeded hf_secret /data/private-user'}], calls); + }); + await dispatcher.dispatch(snapshot); + dispatcher.observe(snapshot); + assert.equal(attempts, 1); + assert.equal(calls.length, 1); + assert.equal(completed.length, 0); + assert.match(errors[0].message, /your personal quota/i); + assert.doesNotMatch(errors[0].message, /hf_secret|\/data\//); + await dispatcher.retry(snapshot); + assert.equal(attempts, 2); +}); + +test('quota verification error is an authentication failure, not exhausted allowance', async () => { + const calls = []; + const {dispatcher, errors} = setup(() => client([{type: 'status', stage: 'error', + message: 'Hugging Face could not verify your personal ZeroGPU quota. Sign in on Hugging Face, then refresh this Space.'}], calls)); + await dispatcher.dispatch(snapshot); + assert.match(errors[0].message, /could not verify your account/i); + assert.match(errors[0].message, /does not mean your quota is exhausted/i); + assert.doesNotMatch(errors[0].message, /wait|resets|quota is insufficient/i); +}); + +test('an unspecified quota error never claims the allowance has been exhausted', async () => { + const calls = []; + const {dispatcher, errors} = setup(() => client([{type: 'status', stage: 'error', message: 'Quota service unavailable'}], calls)); + await dispatcher.dispatch(snapshot); + assert.match(errors[0].message, /quota may still be available/i); + assert.doesNotMatch(errors[0].message, /exhausted|insufficient|reset/i); +}); + +test('a task time limit is not mistaken for insufficient quota just because it mentions quota', async () => { + const calls = []; + const {dispatcher, errors} = setup(() => client([{type: 'data', data: [{status: 'failed', + error: 'The ZeroGPU time limit was reached. Please try again with your available GPU quota.'}]}], calls)); + await dispatcher.dispatch(snapshot); + assert.match(errors[0].message, /task exceeded its GPU time limit/i); + assert.doesNotMatch(errors[0].message, /insufficient|exhausted|reset/i); +}); + +test('account switch while connection loads cannot submit the previous user operation', async () => { + const calls = [], connection = deferred(); + let closed = 0; + const {dispatcher, state, errors, completed} = setup(() => connection.promise); + const pending = dispatcher.dispatch(snapshot); + await tick(); + state.version++; state.identity = {authenticated: true, csrf_token: 'other-user'}; + dispatcher.reset(); + connection.resolve({...client([], calls), close() { closed++; }}); + await pending; await tick(); + assert.equal(calls.length, 0); + assert.equal(errors.length, 0); + assert.equal(completed.length, 0); + assert.equal(closed, 1); +}); + +test('queued callback result is not reported as completed and requires explicit retry', async () => { + const calls = []; + const {dispatcher, completed, errors} = setup(() => client([{type: 'data', data: [{status: 'queued'}]}], calls)); + await dispatcher.dispatch(snapshot); + assert.equal(completed.length, 0); + assert.match(errors[0].message, /busy/i); + dispatcher.observe(snapshot); + assert.equal(calls.length, 1); +}); + +test('terminal server failures surface as failures, not completed Gradio transport', async () => { + const calls = []; + const {dispatcher, completed, errors} = setup(() => client([ + {type: 'data', data: [{status: 'failed', error: 'Your GPU quota is exhausted.'}]}, + {type: 'status', stage: 'complete'} + ], calls)); + await dispatcher.dispatch(snapshot); + assert.equal(completed.length, 0); + assert.match(errors[0].message, /quota/i); +}); + +test('renewed CSRF session creates a new connection for later work', async () => { + const configs = [], calls = []; + const {dispatcher, state} = setup(config => { configs.push(config); return client([], calls); }); + await dispatcher.dispatch(snapshot); + state.identity.csrf_token = 'csrf-2'; + await dispatcher.dispatch({...snapshot, gpu_operation_id: 'b'.repeat(32)}); + assert.deepEqual(configs.map(config => config.headers['X-CSRF-Token']), ['csrf-1', 'csrf-2']); +}); + +test('missing personal quota headers fail before connecting and never claim depleted quota', async () => { + let connections = 0; + const {dispatcher, errors} = setup(() => { connections++; }, {quotaHeaders: async () => ({})}); + await dispatcher.dispatch(snapshot); + assert.equal(connections, 0); + assert.match(errors[0].message, /could not verify your account/); + assert.doesNotMatch(errors[0].message, /quota is insufficient|reset/); +}); + +test('each simultaneous operation protects its own fresh headers from the SDK second handshake', async () => { + const requests = [], optionsSeen = []; + let tokens = 0; + const {dispatcher} = setup(options => { + optionsSeen.push(options); + const instance = { + fetch(url, options) { requests.push({url, headers: new Headers(options.headers)}); }, + submit(endpoint, data, eventData, trigger, allEvents, headers) { + assert.equal(headers['x-gradio-user'], 'app'); + return (async function* () { + await tick(); + await instance.fetch('https://example.hf.space/gradio_api/queue/join', { + headers: {...headers, 'x-zerogpu-token': 'sdk-fallback', 'x-gradio-user': 'api', 'Content-Type': 'application/json'} + }); + await instance.fetch('https://example.hf.space/gradio_api/info', {headers: {}}); + yield {type: 'data', data: [{status: 'complete'}]}; + })(); + } + }; + return instance; + }, {quotaHeaders: async () => ({'x-zerogpu-token': 'token-' + ++tokens})}); + await Promise.all([dispatcher.dispatch(snapshot), dispatcher.dispatch({...snapshot, gpu_operation_id: 'b'.repeat(32)})]); + const gpu = requests.filter(item => item.url.endsWith('/queue/join')); + assert.deepEqual(gpu.map(item => item.headers.get('x-zerogpu-token')).sort(), ['token-3', 'token-4']); + assert.ok(gpu.every(item => item.headers.get('x-gradio-user') === 'app')); + assert.ok(requests.filter(item => item.url.endsWith('/info')).every(item => !item.headers.has('x-zerogpu-token'))); + assert.ok(optionsSeen.every(options => Object.keys(options.headers).join() === 'X-CSRF-Token')); +}); + +function browserEnvironment({hidden = false, token = 'visible-token'} = {}) { + const win = new EventTarget(), doc = new EventTarget(), posts = []; + doc.hidden = hidden; + const parent = {postMessage(message, origin, ports) { + posts.push({message, origin}); + if (message === 'supports-zerogpu-headers') { + const event = new Event('message'); + Object.assign(event, {source: parent, origin: 'https://huggingface.co', data: message}); + queueMicrotask(() => win.dispatchEvent(event)); + } else queueMicrotask(() => ports[0].postMessage(token ? {'X-ZeroGPU-Token': token, 'X-ZeroGPU-UUID': 'test-session'} : {})); + }}; + win.parent = parent; + const context = vm.createContext({window: win, document: doc, MessageChannel, Headers, setTimeout, clearTimeout}); + vm.runInContext(source, context); + return {win, doc, parent, posts, request: context.SupraZeroGPU.requestQuotaHeaders}; +} + +test('quota handshake waits for page visibility and asks HF for support before private headers', async () => { + const env = browserEnvironment({hidden: true}), controller = new AbortController(); + const request = env.request(controller.signal); + await tick(); assert.equal(env.posts.length, 0); + env.doc.hidden = false; env.doc.dispatchEvent(new Event('visibilitychange')); + const headers = await request; + assert.deepEqual(env.posts.map(post => post.message), ['supports-zerogpu-headers', 'zerogpu-headers']); + assert.ok(env.posts.every(post => post.origin === 'https://huggingface.co')); + assert.equal(headers['x-zerogpu-token'], 'visible-token'); + assert.equal(headers['x-zerogpu-uuid'], 'test-session'); + assert.equal(headers['x-gradio-user'], 'app'); +}); + +test('quota handshake rejects absent tokens and cancels hidden-page waits on account reset', async () => { + const missing = browserEnvironment({token: ''}); + await assert.rejects(missing.request(new AbortController().signal), /could not verify/); + const hidden = browserEnvironment({hidden: true}), controller = new AbortController(); + const waiting = hidden.request(controller.signal); + controller.abort(); + await assert.rejects(waiting, {name: 'AbortError'}); + assert.equal(hidden.posts.length, 0); +}); + +test('a token omitted after the page becomes hidden retries when visibility returns', async () => { + const env = browserEnvironment(), send = env.parent.postMessage; + let firstHeaders = true, settled = false; + env.parent.postMessage = (message, origin, ports) => { + if (message === 'zerogpu-headers' && firstHeaders) { + firstHeaders = false; + env.posts.push({message, origin}); + env.doc.hidden = true; + ports[0].postMessage({'X-ZeroGPU-UUID': 'test-session'}); + } else send(message, origin, ports); + }; + const waiting = env.request(new AbortController().signal).then(value => { settled = true; return value; }); + while (firstHeaders) await tick(); + await tick(); + assert.equal(settled, false); + assert.equal(env.posts.length, 2); + env.doc.hidden = false; env.doc.dispatchEvent(new Event('visibilitychange')); + const headers = await waiting; + assert.equal(headers['x-zerogpu-token'], 'visible-token'); + assert.deepEqual(env.posts.map(post => post.message), [ + 'supports-zerogpu-headers', 'zerogpu-headers', 'supports-zerogpu-headers', 'zerogpu-headers' + ]); +}); + +test('final queue send waits for visibility again instead of reusing earlier headers in a hidden page', async () => { + const env = browserEnvironment(), sent = []; + const {dispatcher} = setup(() => { + const instance = { + fetch(url, options) { sent.push(new Headers(options.headers)); }, + submit() { return (async function* () { + env.doc.hidden = true; + await instance.fetch('https://example.hf.space/gradio_api/queue/join', {headers: {'x-zerogpu-token': 'old'}}); + yield {type: 'data', data: [{status: 'complete'}]}; + })(); } + }; + return instance; + }, {quotaHeaders: env.request}, {document: env.doc}); + const pending = dispatcher.dispatch(snapshot); + while (!env.doc.hidden) await tick(); + await tick(); assert.equal(sent.length, 0); + env.doc.hidden = false; env.doc.dispatchEvent(new Event('visibilitychange')); + await pending; + assert.equal(sent.length, 1); + assert.equal(sent[0].get('x-zerogpu-token'), 'visible-token'); + assert.equal(env.posts.filter(post => post.message === 'zerogpu-headers').length, 2); +}); + +test('only the trusted HF parent can acknowledge header support', async () => { + const env = browserEnvironment(), controller = new AbortController(); + env.parent.postMessage = (message, origin) => env.posts.push({message, origin}); + const waiting = env.request(controller.signal); + await tick(); + for (const properties of [ + {source: {}, origin: 'https://huggingface.co'}, + {source: env.parent, origin: 'https://untrusted.example'} + ]) { + const event = new Event('message'); Object.assign(event, {...properties, data: 'supports-zerogpu-headers'}); + env.win.dispatchEvent(event); + } + await tick(); assert.equal(env.posts.length, 1); + controller.abort(); await assert.rejects(waiting, {name: 'AbortError'}); +}); + +test('ZeroGPU workspace is usable before model preparation while GPU execution remains gated', () => { + const context = vm.createContext({}); vm.runInContext(source, context); + const {workspaceReady, modelsReady} = context.SupraZeroGPU; + const pending = {gpu_mode: 'zerogpu', workspace: {stage: 'ready'}, native_build: {stage: 'pending'}, models: {stage: 'pending'}}; + assert.equal(workspaceReady(pending), true); assert.equal(modelsReady(pending), false); + assert.equal(workspaceReady({...pending, gpu_mode: 'dedicated'}), false); + assert.equal(workspaceReady({...pending, workspace: {stage: 'loading'}}), false); + const loaded = {...pending, native_build: {stage: 'compiled'}, models: {stage: 'ready_released'}}; + assert.equal(workspaceReady(loaded), true); assert.equal(modelsReady(loaded), true); + assert.equal(workspaceReady({...loaded, gpu_mode: 'dedicated'}), true); +}); + +test('foreground and generation readiness, failure, and progress use independent model status', () => { + const context = vm.createContext({}); vm.runInContext(source, context); + const {operationReady, preparationError, preparationState} = context.SupraZeroGPU; + const foreground = {gpu_operation_kind: 'foreground'}, generation = {gpu_operation_kind: 'generation'}; + const state = {gpu_mode: 'zerogpu', workspace: {stage: 'ready'}, native_build: {stage: 'building', + startup_progress: {completed: 2, total: 4, unit: 'steps'}}, + foreground: {stage: 'ready_released', startup_progress: {completed: 885e6, total: 885e6, unit: 'bytes'}}, + models: {stage: 'pending', startup_progress: {completed: 0, total: 65.9e9, unit: 'bytes'}}}; + assert.equal(operationReady(state, foreground), true); + assert.equal(operationReady(state, {foreground: {status: 'queued'}}), true); + assert.equal(operationReady(state, generation), false); + assert.equal(preparationState(state, foreground), state.foreground); + assert.equal(preparationState(state, generation), state.native_build); + state.native_build.stage = 'compiled'; + assert.equal(preparationState(state, generation), state.models); + state.models.stage = 'ready_released'; state.foreground.stage = 'failed'; + assert.equal(operationReady(state, foreground), false); + assert.equal(operationReady(state, generation), true); + assert.match(preparationError(state, foreground), /Background-removal/); + assert.equal(preparationError(state, generation), null); + state.workspace.stage = 'starting'; + assert.equal(operationReady(state, generation), false); +}); + +test('preparation percentages use the current phase total with explicit bytes, models, or steps', () => { + const context = vm.createContext({}); vm.runInContext(source, context); + const {formatProgress} = context.SupraZeroGPU; + assert.equal(formatProgress({completed: 442.5e6, total: 885e6, unit: 'bytes'}).text, '50% · 442.5 / 885.0 MB'); + assert.equal(formatProgress({completed: 6.59e9, total: 65.9e9, unit: 'bytes'}).text, '10% · 6.6 / 65.9 GB'); + assert.equal(formatProgress({completed: 2, total: 5, unit: 'models'}).text, '40% · 2 / 5 models'); + assert.equal(formatProgress({completed: 1, total: 4, unit: 'steps'}).text, '25% · 1 / 4 steps'); + assert.equal(formatProgress({label: 'Checking metadata', completed: 0, total: 0, unit: 'steps'}).text, '0%'); + assert.equal(formatProgress({completed: -5, total: 10, unit: 'steps'}).percent, 0); + assert.equal(formatProgress({completed: 15, total: 10, unit: 'steps'}).percent, 100); + assert.equal(formatProgress({completed: Infinity, total: NaN}).percent, 0); +}); + +test('the unmodified official client sends one app marker and preserves fresh HF headers after its own handshake', async () => { + const vendor = fs.readFileSync(path.join(__dirname, '../hf_space/vendor/gradio-client-2.1.0.js'), 'utf8'); + const {Client} = await import('data:text/javascript;base64,' + Buffer.from(vendor).toString('base64')); + const previous = {window: global.window, document: global.document, fetch: global.fetch}; + const requests = []; + let handshakes = 0; + global.window = {location: {hostname: 'example.hf.space', search: ''}, supports_zerogpu_headers: true, + parent: {postMessage(message, origin, ports) { + handshakes++; + assert.equal(message, 'zerogpu-headers'); assert.equal(origin, 'https://huggingface.co'); + ports[0].postMessage({'X-ZeroGPU-Token': 'sdk-second-token', 'X-ZeroGPU-UUID': 'sdk-second-uuid'}); + ports[0].close(); + }}}; + global.document = {}; + global.fetch = async (url, options) => { + requests.push({url, options, headers: new Headers(options.headers)}); + // Finish at input validation: this checks the real queue/join wire request + // without opening a stream, contacting a server, or executing GPU work. + return new Response(JSON.stringify({detail: 'Diagnostic task is not available'}), {status: 422}); + }; + try { + const {dispatcher} = setup(options => { + const instance = new Client('https://example.hf.space', options); + instance.config = {root: 'https://example.hf.space', protocol: 'sse_v3', enable_queue: true, + dependencies: [{id: 0, queue: true, inputs: [1], outputs: [2]}], + components: [{id: 1, type: 'textbox'}, {id: 2, type: 'json'}]}; + instance.api_prefix = '/gradio_api'; + instance.api_info = {named_endpoints: {'/run_gpu_operation': {parameters: [{parameter_name: 'operation_id'}]}}}; + instance.api_map = {run_gpu_operation: 0}; + instance.handle_blob = async (root, data) => data; + return instance; + }); + await dispatcher.dispatch(snapshot); + assert.equal(handshakes, 1); assert.equal(requests.length, 1); + const request = requests[0]; + assert.match(request.url, /\/gradio_api\/queue\/join/); + assert.equal(request.headers.get('x-zerogpu-token'), 'personal-quota-token'); + assert.equal(request.headers.get('x-zerogpu-uuid'), 'session-id'); + assert.equal(request.headers.get('x-gradio-user'), 'app'); + assert.equal(request.headers.get('x-csrf-token'), 'csrf-1'); + assert.equal(request.options.credentials, 'include'); + assert.deepEqual(JSON.parse(request.options.body).data, [id]); + } finally { + if (previous.window === undefined) delete global.window; else global.window = previous.window; + if (previous.document === undefined) delete global.document; else global.document = previous.document; + global.fetch = previous.fetch; + } +}); diff --git a/tests/test_zerogpu_concurrency.py b/tests/test_zerogpu_concurrency.py new file mode 100644 index 0000000000000000000000000000000000000000..c534ba83656170573581c5642ed5d875b234515d --- /dev/null +++ b/tests/test_zerogpu_concurrency.py @@ -0,0 +1,165 @@ +"""Concurrent direct operations use separate state and never spend GPU twice.""" +from concurrent.futures import ThreadPoolExecutor +from io import BytesIO +import threading + +from PIL import Image +import pytest + +from hf_space.history_jobs import Runner +from hf_space.history_store import NotFound, Store + + +@pytest.fixture +def store(tmp_path, monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.delenv('SPACE_ID', raising=False) + monkeypatch.delenv('SPACE_HOST', raising=False) + return Store(tmp_path / 'history', asset_root=tmp_path / 'assets').initialize() + + +def pending(store, owner): + stream = BytesIO() + Image.new('RGB', (16, 16), 'red').save(stream, format='PNG') + upload = store.save_upload(owner, stream.getvalue()) + return store.request_foreground(owner, upload['id'], 'request-' + upload['id'])['gpu_operation_id'] + + +def output(store, job): + color = (1, 2, 3, 255) if job['owner'] == 'alice' else (4, 5, 6, 255) + return store.write_foreground_output( + job, Image.new('RGBA', (16, 16), color), + {'source': 'rmbg2', 'model': 'briaai/RMBG-2.0', 'uncropped': True}) + + +@pytest.mark.parametrize('owners', [('alice', 'bob'), ('alice', 'alice')], + ids=['different-visitors', 'same-visitor-separate-operations']) +def test_operations_enter_handlers_together_without_global_lock(store, owners): + operations = [(owner, pending(store, owner)) for owner in owners] + together = threading.Barrier(2) + entered = set() + guard = threading.Lock() + + def execute(job, update): + with guard: + entered.add(job['operation_id']) + # A serial application queue fails here; both handlers must be active + # before either can finish or return its GPU lease. + together.wait(timeout=5) + update(label='Running ' + job['owner']) + return output(store, job) + + runner = Runner(store, {'foreground': execute}) + with store.execution_lock, ThreadPoolExecutor(max_workers=2) as executor: + tasks = [executor.submit(runner.run_once, operation_id=operation, owner=owner) + for owner, operation in operations] + assert all(task.result(timeout=10) for task in tasks) + assert entered == {operation for _, operation in operations} + for owner, operation in operations: + current = store.get_operation(owner, operation) + assert current['job']['owner'] == owner + assert current['operation']['status'] == 'complete' + + +def test_simultaneous_replays_execute_once_and_still_check_owner(store): + operation = pending(store, 'alice') + entered, release = threading.Event(), threading.Event() + calls = [] + + def execute(job, update): + calls.append(job['operation_id']) + entered.set() + assert release.wait(timeout=5) + return output(store, job) + + runner = Runner(store, {'foreground': execute}) + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(runner.run_once, operation_id=operation, owner='alice') + try: + assert entered.wait(timeout=5) + duplicate = executor.submit(runner.run_once, operation_id=operation, owner='alice') + assert duplicate.result(timeout=2) is False + with pytest.raises(NotFound): + runner.run_once(operation_id=operation, owner='bob') + assert store.get_operation('alice', operation)['operation']['status'] == 'running' + finally: + release.set() + assert first.result(timeout=5) is True + assert runner.run_once(operation_id=operation, owner='alice') is False + assert calls == [operation] + assert store.get_operation('alice', operation)['operation']['status'] == 'complete' + + +def test_concurrent_terminal_failures_retry_only_the_selected_result(store, monkeypatch): + operations = {owner: pending(store, owner) for owner in ('alice', 'bob')} + together = threading.Barrier(2) + calls, writes = {}, {} + guard = threading.Lock() + finish = store.finish_operation + + def execute(job, update): + with guard: + calls[job['owner']] = calls.get(job['owner'], 0) + 1 + together.wait(timeout=5) + return output(store, job) + + def flaky_finish(operation_id, *args, **kwargs): + with guard: + writes[operation_id] = writes.get(operation_id, 0) + 1 + first_write = writes[operation_id] == 1 + if first_write: + raise OSError('temporary storage interruption') + return finish(operation_id, *args, **kwargs) + + monkeypatch.setattr(store, 'finish_operation', flaky_finish) + runner = Runner(store, {'foreground': execute}) + with ThreadPoolExecutor(max_workers=2) as executor: + tasks = [executor.submit(runner.run_once, operation_id=operation, owner=owner) + for owner, operation in operations.items()] + for task in tasks: + with pytest.raises(OSError, match='storage interruption'): + task.result(timeout=10) + with pytest.raises(NotFound): + runner.run_once(operation_id=operations['alice'], owner='bob') + assert runner.run_once() is False # pending writes never revive a background FIFO + assert runner.run_once(operation_id=operations['alice'], owner='alice') is True + assert store.get_operation('alice', operations['alice'])['operation']['status'] == 'complete' + assert store.get_operation('bob', operations['bob'])['operation']['status'] == 'running' + assert runner.run_once(operation_id=operations['bob'], owner='bob') is True + assert store.get_operation('bob', operations['bob'])['operation']['status'] == 'complete' + assert calls == {'alice': 1, 'bob': 1} + + +def test_zero_mode_cannot_start_an_anonymous_background_worker(store): + operation = pending(store, 'alice') + runner = Runner(store, {'foreground': lambda *_: pytest.fail('anonymous GPU work')}) + assert runner.start() is None + assert runner.run_once() is False + assert store.get_operation('alice', operation)['operation']['status'] == 'queued' + + +def test_paid_mode_keeps_serial_execution(store, monkeypatch): + alice, bob = pending(store, 'alice'), pending(store, 'bob') + monkeypatch.delenv('SPACES_ZERO_GPU') + entered, release = threading.Event(), threading.Event() + calls = [] + + def execute(job, update): + calls.append(job['owner']) + if job['owner'] == 'alice': + entered.set() + assert release.wait(timeout=5) + return output(store, job) + + runner = Runner(store, {'foreground': execute}) + with ThreadPoolExecutor(max_workers=1) as executor: + first = executor.submit(runner.run_once, operation_id=alice, owner='alice') + try: + assert entered.wait(timeout=5) + assert runner.run_once(operation_id=bob, owner='bob') is False + assert store.get_operation('bob', bob)['operation']['status'] == 'queued' + finally: + release.set() + assert first.result(timeout=5) is True + assert runner.run_once(operation_id=bob, owner='bob') is True + assert calls == ['alice', 'bob'] diff --git a/tests/test_zerogpu_jobs.py b/tests/test_zerogpu_jobs.py new file mode 100644 index 0000000000000000000000000000000000000000..410179ec9c1b6d3381e614d49ee4cdb0bb04293a --- /dev/null +++ b/tests/test_zerogpu_jobs.py @@ -0,0 +1,129 @@ +import asyncio +from io import BytesIO +from types import SimpleNamespace + +import gradio as gr +from gradio.context import LocalContext +from PIL import Image +import pytest +from starlette.requests import Request + +from hf_space.history_auth import Identity, csrf_token +from hf_space.history_jobs import Runner +from hf_space.history_store import Store, NotFound +from hf_space.zerogpu_jobs import register + + +@pytest.fixture +def store(tmp_path, monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.delenv('SPACE_ID', raising=False) + monkeypatch.delenv('SPACE_HOST', raising=False) + return Store(tmp_path / 'history', asset_root=tmp_path / 'assets').initialize() + + +def pending(store, owner): + stream = BytesIO() + Image.new('RGB', (16, 16), 'red').save(stream, format='PNG') + upload = store.save_upload(owner, stream.getvalue()) + upload = store.request_foreground(owner, upload['id'], 'request-' + upload['id']) + return upload['gpu_operation_id'] + + +def request(owner, *, quota=True, csrf=True, authenticated=True): + identity = Identity(owner, owner, 'session-' + owner) + headers = [(b'host', b'testserver'), (b'origin', b'https://testserver')] + if quota: + headers.append((b'x-ip-token', ('quota-' + owner).encode())) + if csrf: + headers.append((b'x-csrf-token', csrf_token(identity).encode())) + scope = dict(type='http', method='POST', scheme='https', server=('testserver', 443), + path='/gradio_api/queue/join', root_path='', query_string=b'', headers=headers, + client=('127.0.0.1', 1), session={}) + if authenticated: + scope['supra_history_app_identity'] = identity + return gr.Request(request=Request(scope), session_hash='test-' + owner) + + +def make_app(store, handler): + runner = Runner(store, {'foreground': handler}) + service = SimpleNamespace(store=store, runner=runner, ready=lambda kind: True) + with gr.Blocks() as demo: + pass + callback = register(demo, service) + return demo, callback, runner + + +def process(demo, operation, req): + result = asyncio.run(demo.process_api(0, [operation], request=req))['data'][0] + return result.model_dump() if hasattr(result, 'model_dump') else result + + +def test_two_accounts_keep_own_context_and_exact_operation(store): + first, second = pending(store, 'alice'), pending(store, 'bob') + seen = [] + def execute(job, update): + seen.append((job['owner'], LocalContext.request.get().headers['x-ip-token'])) + return store.write_foreground_output(job, Image.new('RGBA', (16,16), (1,2,3,255)), + {'source': 'rmbg2', 'model': 'briaai/RMBG-2.0', 'uncropped': True}) + demo, _, runner = make_app(store, execute) + assert runner.run_once() is False # no anonymous background fallback + # Bob's request must not execute the earlier Alice operation with Bob's quota. + assert process(demo, second, request('bob'))['status'] == 'complete' + assert store.get_operation('alice', first)['operation']['status'] == 'queued' + assert process(demo, first, request('alice'))['status'] == 'complete' + assert seen == [('bob', 'quota-bob'), ('alice', 'quota-alice')] + assert LocalContext.request.get(None) is None + # A repeated request must never spend GPU twice. + assert process(demo, first, request('alice'))['status'] == 'complete' + assert len(seen) == 2 + + +@pytest.mark.parametrize('kwargs', [dict(quota=False), dict(csrf=False), dict(authenticated=False)]) +def test_rejects_missing_auth_before_execution(store, kwargs): + operation = pending(store, 'alice') + demo, callback, _ = make_app(store, lambda *args: pytest.fail('GPU must not run')) + with pytest.raises(gr.Error): + callback(operation, request('alice', **kwargs)) + assert store.get_operation('alice', operation)['operation']['status'] == 'queued' + + +def test_cross_account_operation_hidden(store): + operation = pending(store, 'alice') + _, callback, _ = make_app(store, lambda *args: pytest.fail('GPU must not run')) + with pytest.raises(gr.Error, match='not available'): + callback(operation, request('bob')) + with pytest.raises(NotFound): + store.get_operation('bob', operation) + + +def test_quota_failure_saved_without_token_and_no_retry(store): + operation = pending(store, 'alice') + calls = [] + def failure(*args): + calls.append(True) + raise RuntimeError('ZeroGPU quota exceeded; sensitive-token-placeholder') + demo, _, _ = make_app(store, failure) + result = process(demo, operation, request('alice')) + assert result['status'] == 'failed' + assert 'quota' in result['error'] + assert 'sensitive-token' not in result['error'] + process(demo, operation, request('alice')) + assert len(calls) == 1 + + +def test_restart_restores_queued_work_without_credentials(store): + operation = pending(store, 'alice') + restored = Store(store.root, asset_root=store.asset_root).initialize() + assert restored.get_operation('alice', operation)['operation']['status'] == 'queued' + assert Runner(restored, {}).run_once() is False + for path in store.root.rglob('*.json'): + assert 'x-ip-token' not in path.read_text() + + +def test_nonzero_mode_retains_existing_background_runner(store, monkeypatch): + operation = pending(store, 'alice') + monkeypatch.delenv('SPACES_ZERO_GPU') + runner = Runner(store, {'foreground': lambda *args: (_ for _ in ()).throw(RuntimeError('test'))}) + assert runner.run_once() is True + assert store.get_operation('alice', operation)['operation']['status'] == 'failed' diff --git a/tests/test_zerogpu_lazy.py b/tests/test_zerogpu_lazy.py new file mode 100644 index 0000000000000000000000000000000000000000..17062a33f106dd190cc750e4ba9c2b0bc88124c5 --- /dev/null +++ b/tests/test_zerogpu_lazy.py @@ -0,0 +1,114 @@ +"""Browsing saved work is independent of optional, owner-requested CPU loading.""" +from io import BytesIO +from types import SimpleNamespace +from unittest.mock import Mock +import time + +from fastapi import FastAPI +from PIL import Image +import pytest +from starlette.testclient import TestClient + +from hf_space.history_auth import Identity, csrf_token +from hf_space.history_bridge import AppSessions, AppSessionMiddleware, COOKIE_NAME +from hf_space.history_routes import install_routes +from hf_space.history_store import Store + + +def test_public_status_exposes_runtime_epoch_without_private_diagnostics(monkeypatch): + from hf_space import history_service + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + models = {'stage': 'pending', 'epoch': 'a' * 32, + 'error': 'private model traceback', 'config': '/private/model/config.json'} + monkeypatch.setattr(history_service.release_backend, 'model_status', lambda: dict(models)) + monkeypatch.setattr(history_service.release_backend, 'foreground_status', lambda: { + 'stage': 'pending', 'error': 'private foreground traceback'}) + monkeypatch.setattr(history_service.native_build, 'status', lambda: { + 'stage': 'pending', 'error': 'private native traceback', 'resources': {'internal': 'private'}}) + first = history_service.deployment_status() + assert first['runtime_id'] == 'a' * 32 + assert 'private' not in str(first) + assert first['models'] == {'stage': 'pending'} + models['epoch'] = 'b' * 32 + second = history_service.deployment_status() + assert second['runtime_id'] == 'b' * 32 + assert second['runtime_id'] != first['runtime_id'] + assert 'private' not in str(second) + + +@pytest.fixture +def lazy_app(tmp_path, monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + for name in ('SPACE_ID', 'SPACE_HOST', 'OAUTH_CLIENT_ID', 'OAUTH_CLIENT_SECRET'): + monkeypatch.delenv(name, raising=False) + store = Store(tmp_path/'history', asset_root=tmp_path/'assets').initialize() + sessions = AppSessions('test-only-lazy-session') + prepare = Mock() + example = tmp_path/'example.png' + Image.new('RGB', (16, 16), 'red').save(example) + web = FastAPI() + install_routes(web, store, deployment_status=lambda: { + 'gpu_mode': 'zerogpu', 'native_build': {'stage': 'pending'}, + 'models': {'stage': 'pending'}}, prepare_models=prepare, + examples_provider=lambda: [], example_file_provider=lambda name, kind: example, + private_cache_root=tmp_path/'private-cache') + web.add_middleware(AppSessionMiddleware, sessions=sessions) + with TestClient(web, base_url='https://testserver') as client: + yield SimpleNamespace(client=client, store=store, prepare=prepare, sessions=sessions) + + +def headers(app, owner='alice'): + claims = app.sessions.claims(Identity(owner, owner, 'test-oauth'), time.time()+600) + identity = app.sessions.validate(claims) + return {'Origin': 'https://testserver', 'Sec-Fetch-Site': 'same-origin', + 'X-CSRF-Token': csrf_token(identity), + 'Cookie': COOKIE_NAME+'='+app.sessions.serializer.dumps(claims)} + + +def upload(app): + stream = BytesIO() + Image.new('RGB', (16, 16), 'red').save(stream, format='PNG') + response = app.client.post('/api/history/upload', headers=headers(app), + files={'file': ('image.png', stream.getvalue(), 'image/png')}) + assert response.status_code == 200, response.text + return response.json() + + +def test_browsing_and_upload_admission_do_not_start_models(lazy_app): + app = lazy_app + for path in ('/api/history/status', '/api/history/examples', + '/api/history/examples/example/input', '/api/history/jobs'): + response = app.client.get(path, headers=headers(app)) + assert response.status_code == 200, response.text + value = upload(app) + assert value['gpu_operation_id'] + assert app.client.get('/api/history/jobs', headers=headers(app)).status_code == 200 + app.prepare.assert_not_called() + + +def test_only_owner_explicit_prepare_starts_cpu_loading(lazy_app): + app = lazy_app + value = upload(app) + data = {'operation_id': value['gpu_operation_id']} + response = app.client.post('/api/history/gpu/prepare', headers=headers(app, 'bob'), json=data) + assert response.status_code == 404 + missing_csrf = headers(app) + del missing_csrf['X-CSRF-Token'] + assert app.client.post('/api/history/gpu/prepare', headers=missing_csrf, json=data).status_code == 403 + app.prepare.assert_not_called() + response = app.client.post('/api/history/gpu/prepare', headers=headers(app), json=data) + assert response.status_code == 200, response.text + assert response.json()['status'] == 'queued' + app.prepare.assert_called_once_with('foreground') + + +def test_finished_task_does_not_reload_models(lazy_app): + app = lazy_app + value = upload(app) + operation = value['gpu_operation_id'] + app.store.claim_operation(operation) + app.store.finish_operation(operation, error='test operation ended') + response = app.client.post('/api/history/gpu/prepare', headers=headers(app), + json={'operation_id': operation}) + assert response.status_code == 200 + app.prepare.assert_not_called() diff --git a/tests/test_zerogpu_queue.py b/tests/test_zerogpu_queue.py new file mode 100644 index 0000000000000000000000000000000000000000..ba3ba07915583ed218f92d6c095c1d177edfc3e2 --- /dev/null +++ b/tests/test_zerogpu_queue.py @@ -0,0 +1,185 @@ +"""Real ASGI/Gradio queue coverage; inference is a local CPU-only stub.""" +import base64 +import json +import os +import time +import uuid +from io import BytesIO +from types import SimpleNamespace + +os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False') + +import gradio as gr +from fastapi import FastAPI +from gradio.context import LocalContext +from PIL import Image +import pytest +from starlette.testclient import TestClient + +from hf_space.history_auth import Identity, csrf_token +from hf_space.history_bridge import AppSessionMiddleware, AppSessions, COOKIE_NAME +from hf_space.history_jobs import Runner +from hf_space.history_store import Store +from hf_space.zerogpu_jobs import register + + +@pytest.fixture +def queued_app(tmp_path, monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.delenv('SPACE_ID', raising=False) + monkeypatch.delenv('SPACE_HOST', raising=False) + monkeypatch.delenv('OAUTH_CLIENT_ID', raising=False) + monkeypatch.delenv('OAUTH_CLIENT_SECRET', raising=False) + store = Store(tmp_path / 'history', asset_root=tmp_path / 'assets').initialize() + sessions = AppSessions('test-only-queue-cookie-secret') + seen = [] + + def execute(job, update): + request = LocalContext.request.get(None) + assert request is not None + assert request.request is not None + identity = request.request.scope['supra_history_app_identity'] + assert identity.owner == job['owner'] + assert request.request.url.path == '/gradio_api/queue/join' + seen.append((job['owner'], request.headers['x-ip-token'])) + return store.write_foreground_output( + job, Image.new('RGBA', (16, 16), (1, 2, 3, 255)), + {'source': 'rmbg2', 'model': 'briaai/RMBG-2.0', 'uncropped': True}) + + runner = Runner(store, {'foreground': execute}) + service = SimpleNamespace(store=store, runner=runner, ready=lambda kind: True) + with gr.Blocks() as demo: + pass + register(demo, service) + demo.queue(api_open=False, max_size=32) + web = gr.mount_gradio_app(FastAPI(), demo, path='/', ssr_mode=False) + gradio_app = next(route.app for route in web.routes if getattr(route, 'path', None) == '') + gradio_app.add_middleware(AppSessionMiddleware, sessions=sessions) + dependency = next(item for item in demo.config['dependencies'] + if item.get('api_name') == 'run_gpu_operation') + with TestClient(web, base_url='https://testserver') as client: + yield SimpleNamespace(client=client, demo=demo, store=store, sessions=sessions, + seen=seen, fn_index=dependency['id']) + + +def _pending(app, owner): + stream = BytesIO() + Image.new('RGB', (16, 16), 'red').save(stream, format='PNG') + upload = app.store.save_upload(owner, stream.getvalue()) + upload = app.store.request_foreground(owner, upload['id'], 'request-' + upload['id']) + return upload['gpu_operation_id'] + + +def _headers(app, owner, *, cookie=True, quota=True, csrf=True): + claims = app.sessions.claims(Identity(owner, owner, 'oauth-session'), time.time() + 600) + identity = app.sessions.validate(claims) + headers = {'Origin': 'https://testserver', 'Sec-Fetch-Site': 'same-origin'} + if cookie: + headers['Cookie'] = COOKIE_NAME + '=' + app.sessions.serializer.dumps(claims) + if csrf: + headers['X-CSRF-Token'] = csrf_token(identity) + if quota: + headers['X-IP-Token'] = 'test-quota-' + owner + return headers + + +def _join(app, operation_id, headers): + session_hash = uuid.uuid4().hex + response = app.client.post('/gradio_api/queue/join', + json={'data': [operation_id], 'fn_index': app.fn_index, + 'session_hash': session_hash}, headers=headers) + assert response.status_code == 200, response.text + assert response.json()['event_id'] + return session_hash + + +def _completion(app, session_hash, headers): + with app.client.stream('GET', '/gradio_api/queue/data', + params={'session_hash': session_hash}, headers=headers) as response: + assert response.status_code == 200 + for line in response.iter_lines(): + if line.startswith('data: '): + event = json.loads(line[6:]) + if event.get('msg') == 'process_completed': + return event + pytest.fail('Gradio queue ended without process_completed') + + +def test_real_queue_keeps_each_signed_session_and_quota_and_deduplicates(queued_app): + app = queued_app + # The browser client reads default /info (or the equivalent root-page + # global), not ?all_endpoints=true. Keyword submission needs this schema. + response = app.client.get('/gradio_api/info') + assert response.status_code == 200 + endpoint = response.json()['named_endpoints']['/run_gpu_operation'] + assert [parameter['parameter_name'] for parameter in endpoint['parameters']] == ['operation_id'] + alice, bob = _pending(app, 'alice'), _pending(app, 'bob') + alice_headers, bob_headers = _headers(app, 'alice'), _headers(app, 'bob') + # The older Alice operation must not run under Bob's quota. Both queue + # admissions happen before reading SSE so they overlap in the same worker. + bob_session = _join(app, bob, bob_headers) + alice_session = _join(app, alice, alice_headers) + for session_hash, headers in ((bob_session, bob_headers), (alice_session, alice_headers)): + event = _completion(app, session_hash, headers) + assert event['success'] is True, event + assert event['output']['data'][0]['status'] == 'complete' + # Independent operations can finish in either order; quota identity must + # remain attached to its own signed session throughout both executions. + assert sorted(app.seen) == [('alice', 'test-quota-alice'), ('bob', 'test-quota-bob')] + + replay = _completion(app, _join(app, alice, alice_headers), alice_headers) + assert replay['success'] is True + assert replay['output']['data'][0]['status'] == 'complete' + assert len(app.seen) == 2 + assert LocalContext.request.get(None) is None + + +@pytest.mark.parametrize('missing', ['cookie', 'csrf', 'quota']) +def test_real_queue_rejects_missing_identity_csrf_or_quota(queued_app, missing): + app = queued_app + operation_id = _pending(app, 'alice') + headers = _headers(app, 'alice', **{missing: False}) + event = _completion(app, _join(app, operation_id, headers), headers) + assert event['success'] is False + assert not app.seen + assert app.store.get_operation('alice', operation_id)['operation']['status'] == 'queued' + + +def test_real_queue_rejects_other_accounts_operation(queued_app): + app = queued_app + operation_id = _pending(app, 'alice') + headers = _headers(app, 'bob') + event = _completion(app, _join(app, operation_id, headers), headers) + assert event['success'] is False + assert 'not available' in event['output']['error'] + assert not app.seen + assert app.store.get_operation('alice', operation_id)['operation']['status'] == 'queued' + + +def test_real_queue_rejects_ip_quota_fallback_without_poisoning_retry(queued_app): + app = queued_app + operation_id = _pending(app, 'alice') + headers = _headers(app, 'alice') + error = 'sensitive-provider-error-placeholder' + payload = base64.urlsafe_b64encode(json.dumps({'error': error}).encode()).decode().rstrip('=') + headers['X-IP-Token'] = 'test-header.' + payload + '.test-signature' + + event = _completion(app, _join(app, operation_id, headers), headers) + assert event['success'] is False + assert 'could not verify your personal ZeroGPU quota' in event['output']['error'] + assert error not in json.dumps(event) + assert not app.seen + assert app.store.get_operation('alice', operation_id)['operation']['status'] == 'queued' + + # A new valid request can retry the same durable job. Decoded claims must + # never replace the signed app identity or mutate the scheduler token. + payload = base64.urlsafe_b64encode(json.dumps({'sub': 'bob', 'error': None}).encode()).decode().rstrip('=') + headers['X-IP-Token'] = 'test-header.' + payload + '.test-signature' + event = _completion(app, _join(app, operation_id, headers), headers) + assert event['success'] is True + assert event['output']['data'][0]['status'] == 'complete' + assert app.seen == [('alice', headers['X-IP-Token'])] + for path in app.store.root.rglob('*.json'): + saved = path.read_text() + assert error not in saved + assert headers['X-IP-Token'] not in saved diff --git a/tests/test_zerogpu_runtime.py b/tests/test_zerogpu_runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..de87c0dbbe4421240e3cb510063d4cb6704bbc1a --- /dev/null +++ b/tests/test_zerogpu_runtime.py @@ -0,0 +1,273 @@ +"""CPU-only checks for the boundary between HTTP state and GPU leases.""" +from pathlib import Path +from concurrent.futures import ThreadPoolExecutor +import importlib.util +import sys +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from hf_space import native_build, release_backend, zerogpu_runtime + + +@pytest.mark.parametrize('zero_gpu', [True, False]) +def test_app_import_and_zero_startup_do_not_prepare_models(monkeypatch, zero_gpu): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true' if zero_gpu else 'false') + monkeypatch.setitem(sys.modules, 'spaces', SimpleNamespace()) + native_start, prepare, handshake = Mock(), Mock(), Mock() + monkeypatch.setattr(native_build, 'start', native_start) + monkeypatch.setattr(release_backend, 'start_preparation', prepare) + monkeypatch.setattr(zerogpu_runtime, 'start_mounted_app', handshake) + source = Path(__file__).resolve().parents[1] / 'app.py' + spec = importlib.util.spec_from_file_location('_test_demand_loaded_app', source) + app = importlib.util.module_from_spec(spec) + spec.loader.exec_module(app) + native_start.assert_not_called() + prepare.assert_not_called() + handshake.assert_not_called() + service = SimpleNamespace(start=Mock()) + web = object() + monkeypatch.setattr(app, 'Service', lambda: service) + monkeypatch.setattr(app, 'create_app', Mock(return_value=web)) + monkeypatch.setattr(app.uvicorn, 'run', Mock()) + app.main() + service.start.assert_called_once_with() + app.create_app.assert_called_once_with(service) + assert app.uvicorn.run.call_args.args == (web,) + assert handshake.call_count == int(zero_gpu) + assert native_start.call_count == int(not zero_gpu) + assert prepare.call_count == int(not zero_gpu) + + +def test_zero_app_keeps_gradio_transport_without_global_admission_cap(monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setitem(sys.modules, 'spaces', SimpleNamespace()) + source = Path(__file__).resolve().parents[1] / 'app.py' + spec = importlib.util.spec_from_file_location('_test_unlimited_zero_app', source) + app = importlib.util.module_from_spec(spec) + spec.loader.exec_module(app) + # Exercise the real pinned Gradio behavior: queue(None) normally becomes + # queue(1) on ZeroGPU, so merely changing the public argument is insufficient. + monkeypatch.setattr(app.gr.utils, 'is_zero_gpu_space', lambda: True) + with app.gr.Blocks() as demo: + pass + monkeypatch.setattr(app, 'build_ui', lambda _: demo) + install = Mock() + monkeypatch.setattr(app, 'install_routes', install) + service = SimpleNamespace(store=object(), examples=Mock(), example_file=Mock(), + prepare_example=Mock(), preview=Mock(), prepare_models=Mock(), + private_files=object()) + app.create_app(service) + assert demo._queue.max_size is None + assert demo.api_open is False + assert install.call_args.kwargs['prepare_models'] is service.prepare_models + service.prepare_models.assert_not_called() + + +def test_concurrent_requests_share_one_preparation_controller(monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setattr(release_backend, '_started', False) + monkeypatch.setattr(release_backend, '_worker', None) + entered, release = threading.Event(), threading.Event() + calls = [] + + def prepare(): + calls.append(threading.get_ident()) + entered.set() + assert release.wait(timeout=5) + + monkeypatch.setattr(release_backend, 'prepare_models', prepare) + together = threading.Barrier(2) + def submit(): + together.wait(timeout=5) + return release_backend.start_preparation() + try: + with ThreadPoolExecutor(max_workers=2) as executor: + tasks = [executor.submit(submit) for _ in range(2)] + workers = [task.result(timeout=5) for task in tasks] + assert entered.wait(timeout=5) + assert workers[0] is workers[1] + assert len(calls) == 1 + finally: + release.set() + if release_backend._worker is not None: + release_backend._worker.join(timeout=5) + + +def test_zero_resident_calls_can_enter_separate_gpu_leases_together(monkeypatch, tmp_path): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setattr(release_backend, 'model_status', lambda: {'stage': 'ready_released'}) + together = threading.Barrier(2) + + class PaidOnlyLock: + def __enter__(self): + pytest.fail('ZeroGPU must not acquire the dedicated worker lock') + + def __exit__(self, *args): + pass + + def generate(path, seed, prepared): + together.wait(timeout=5) + return seed + + monkeypatch.setattr(release_backend, '_generation_lock', PaidOnlyLock()) + monkeypatch.setattr(zerogpu_runtime, 'generate', generate) + with ThreadPoolExecutor(max_workers=2) as executor: + tasks = [executor.submit(release_backend.run_resident, tmp_path / str(seed), seed, 180, None) + for seed in (11, 22)] + assert [task.result(timeout=10) for task in tasks] == [11, 22] + + +def test_cpu_preload_keeps_the_sdk_initialized_interop_pool(monkeypatch): + from hf_space.resident_pipeline import _configure_preload_threads + torch = SimpleNamespace(set_num_threads=Mock(), get_num_interop_threads=Mock(return_value=8), + set_num_interop_threads=Mock(side_effect=RuntimeError('interop already configured by spaces'))) + _configure_preload_threads(torch, cpu_only=True) + torch.set_num_threads.assert_called_once_with(1) + torch.set_num_interop_threads.assert_not_called() + + +def test_dedicated_gpu_worker_still_configures_its_fresh_interop_pool(): + from hf_space.resident_pipeline import _configure_preload_threads + torch = SimpleNamespace(set_num_threads=Mock(), get_num_interop_threads=Mock(return_value=8), + set_num_interop_threads=Mock()) + _configure_preload_threads(torch, cpu_only=False) + torch.set_num_threads.assert_called_once_with(4) + torch.set_num_interop_threads.assert_called_once_with(1) + + +def test_mounted_startup_only_runs_once_and_only_on_zero_gpu(monkeypatch): + startup = Mock() + monkeypatch.setitem(sys.modules, 'spaces.zero', SimpleNamespace(startup=startup)) + monkeypatch.setattr(zerogpu_runtime, '_mounted_started', False) + monkeypatch.setenv('SPACES_ZERO_GPU', 'false') + zerogpu_runtime.start_mounted_app() + startup.assert_not_called() + assert zerogpu_runtime._mounted_started is False + + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + zerogpu_runtime.start_mounted_app() + zerogpu_runtime.start_mounted_app() + startup.assert_called_once_with() + assert zerogpu_runtime._mounted_started is True + + +def test_failed_mounted_startup_can_be_retried(monkeypatch): + startup = Mock(side_effect=[RuntimeError('startup report failed'), None]) + monkeypatch.setitem(sys.modules, 'spaces.zero', SimpleNamespace(startup=startup)) + monkeypatch.setattr(zerogpu_runtime, '_mounted_started', False) + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + with pytest.raises(RuntimeError, match='startup report failed'): + zerogpu_runtime.start_mounted_app() + assert zerogpu_runtime._mounted_started is False + zerogpu_runtime.start_mounted_app() + assert startup.call_count == 2 + assert zerogpu_runtime._mounted_started is True + + +def test_zero_boot_preloads_host_models_without_starting_cuda_process(monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + preload = Mock(return_value={'loaded_models': 18}) + update = Mock() + monkeypatch.setattr(zerogpu_runtime, 'preload', preload) + monkeypatch.setattr(release_backend, 'update', update) + monkeypatch.setattr(release_backend.subprocess, 'Popen', Mock(side_effect=AssertionError('GPU process at boot'))) + monkeypatch.setattr(release_backend.threading, 'Thread', Mock(side_effect=AssertionError('GPU monitor at boot'))) + release_backend.launch_runtime('cpu-config.json') + release_backend._start_runtime_monitor() + assert preload.call_args.args[0] == 'cpu-config.json' + assert update.call_args.kwargs['model_loading'] == 'cpu_pool_with_per_user_zerogpu' + + +def test_preload_only_publishes_a_complete_cpu_pool(monkeypatch): + pool = SimpleNamespace(sealed=True, summary=lambda: {'loaded_models': 18}) + resident = SimpleNamespace(pool=pool, prepare=Mock()) + monkeypatch.setattr(zerogpu_runtime, '_resident', None) + notify = Mock() + assert zerogpu_runtime.preload('config', notify, factory=lambda _: resident) == {'loaded_models': 18} + resident.prepare.assert_called_once_with(notify, cpu_only=True) + assert zerogpu_runtime._resident is resident + pool.sealed = False + with pytest.raises(RuntimeError, match='not sealed'): + zerogpu_runtime.preload('config', notify, factory=lambda _: resident) + + +@pytest.mark.parametrize('operation', ['generate', 'remove_background']) +def test_zero_dispatch_only_sends_serializable_job_arguments(monkeypatch, tmp_path, operation): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setattr(release_backend, 'model_status', lambda: {'stage': 'ready_released'}) + monkeypatch.setattr(release_backend, 'foreground_status', lambda: {'stage': 'ready_released'}) + generate, foreground = Mock(return_value={'ok': True}), Mock(return_value={'ok': True}) + monkeypatch.setattr(zerogpu_runtime, 'generate', generate) + monkeypatch.setattr(zerogpu_runtime, 'foreground', foreground) + monkeypatch.setattr(release_backend, '_runtime_process', None) + assert release_backend.run_resident(tmp_path, 42, 180, None, operation=operation, + foreground_prepared=True) == {'ok': True} + if operation == 'generate': + generate.assert_called_once_with(str(tmp_path), 42, True) + foreground.assert_not_called() + else: + foreground.assert_called_once_with(str(tmp_path)) + generate.assert_not_called() + + +def test_failed_gpu_execution_restores_the_inherited_model_pool(monkeypatch, tmp_path): + from hf_space import runtime_worker + restore = Mock() + resident = SimpleNamespace(pool=SimpleNamespace(restore_all=restore)) + monkeypatch.setattr(zerogpu_runtime, '_resident', resident) + monkeypatch.setattr(runtime_worker, 'execute', Mock(side_effect=ValueError('inference failure'))) + with pytest.raises(ValueError, match='inference failure'): + zerogpu_runtime._run(str(tmp_path), 42) + restore.assert_called_once_with() + assert 'inference failure' in (tmp_path/'inference.log').read_text() + + +def test_zero_native_cache_is_separate_and_uses_advertised_architecture(monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setenv('ZEROGPU_CUDA_CAPABILITY_MAJOR', '12') + monkeypatch.setenv('ZEROGPU_CUDA_CAPABILITY_MINOR', '0') + assert native_build.cuda_architectures() == '12.0+PTX' + assert native_build.wheel_cache_key() != 'torch280-cu128-cp310-sm89-v1' + assert native_build.runtime_environment()['TORCH_CUDA_ARCH_LIST'] == '12.0+PTX' + monkeypatch.delenv('ZEROGPU_CUDA_CAPABILITY_MAJOR') + assert native_build.cuda_architectures() == '9.0;12.0+PTX' + monkeypatch.setenv('SPACES_ZERO_GPU', 'false') + assert native_build.cuda_architectures() == '8.9+PTX' + + +def test_zero_reprojection_reuses_gpu_process_and_restores_argv(monkeypatch, tmp_path): + import importlib + source = Path(__file__).resolve().parents[1]/'supratex/inference/reprojection/__init__.py' + spec = importlib.util.spec_from_file_location('_test_zero_reprojection', source, + submodule_search_locations=[str(source.parent)]) + reprojection = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, reprojection) + spec.loader.exec_module(reprojection) + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + original = sys.argv + captured = [] + stage = SimpleNamespace(main=lambda: captured.append(tuple(sys.argv))) + monkeypatch.setattr(importlib, 'import_module', lambda name: stage) + monkeypatch.setattr(reprojection.subprocess, 'Popen', Mock(side_effect=AssertionError('CUDA subprocess'))) + clear = Mock() + monkeypatch.setitem(sys.modules, 'torch', SimpleNamespace(cuda=SimpleNamespace(empty_cache=clear))) + reprojection._command('_project', ['--output', 'test-output'], tmp_path/'project.log') + assert captured == [('_project', '--output', 'test-output')] + assert sys.argv is original + clear.assert_called_once_with() + + +def test_zero_recovery_formats_preview_in_parent(monkeypatch): + from hf_space import ovoxel_results, preview_recovery + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setattr(ovoxel_results, 'directory', lambda _: Path('missing-test-asset')) + gpu = Mock(return_value='a'*32) + html = Mock(return_value='') + monkeypatch.setattr(zerogpu_runtime, 'recover', gpu) + monkeypatch.setattr(ovoxel_results, 'preview_html', html) + assert preview_recovery.recover('a'*32) == ('', 'a'*32) + gpu.assert_called_once_with('a'*32) + html.assert_called_once_with('a'*32) diff --git a/tests/test_zerogpu_runtime_foreground.py b/tests/test_zerogpu_runtime_foreground.py new file mode 100644 index 0000000000000000000000000000000000000000..69611e72224e1ee55a861b18e6e5a7b6c899ce44 --- /dev/null +++ b/tests/test_zerogpu_runtime_foreground.py @@ -0,0 +1,113 @@ +"""A cutout prepares only RMBG and can finish before 3D models/native builds.""" +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +from hf_space import release_backend, zerogpu_runtime + + +def test_foreground_downloads_only_the_pinned_verified_rmbg_group(monkeypatch): + from hf_space import download_models, weight_cache + inventory = Mock(return_value=[{'size': 10}]) + register = Mock() + mirror = Mock(side_effect=lambda source, destination, progress: progress(source/'model.safetensors', 5, 10)) + update = Mock() + monkeypatch.setattr(release_backend, '_update_foreground', update) + monkeypatch.setattr(download_models, 'source_inventory', inventory) + monkeypatch.setattr(weight_cache, 'register_sources', register) + monkeypatch.setattr(weight_cache, '_mirror', mirror) + path = release_backend._stage_foreground_models() + operation, root = inventory.call_args.args + assert operation['repo_id'] == 'briaai/RMBG-2.0' + assert operation['revision'] == '5df4c9c76d8170882c34f6986e848ee07fd0ba43' + assert operation['destination'] == 'external/rmbg2' + assert 'model.safetensors' in operation['allow_patterns'] + assert len(inventory.call_args_list) == 1 + register.assert_called_once() + assert register.call_args.args == ([{'size': 10}],) + assert mirror.call_args.args[:2] == (root/'external/rmbg2', path) + assert update.call_args.kwargs['startup_progress'] == { + 'label':'Preparing background-removal model', 'completed':5, 'total':10, 'unit':'bytes'} + + +def test_foreground_resident_constructs_only_one_cpu_model(monkeypatch): + from hf_space import resident_pipeline + model = SimpleNamespace(eval=Mock()) + model.eval.return_value = model + load = Mock(return_value=model) + pool = SimpleNamespace(add=Mock(), sealed=False) + monkeypatch.setitem(sys.modules, 'torch', SimpleNamespace(set_num_threads=Mock())) + monkeypatch.setitem(sys.modules, 'transformers', SimpleNamespace( + AutoModelForImageSegmentation=SimpleNamespace(from_pretrained=load))) + monkeypatch.setattr(resident_pipeline, 'ModelPool', lambda **_: pool) + resident = resident_pipeline.ForegroundResident('verified-rmbg-path') + resident.prepare(Mock()) + load.assert_called_once_with('verified-rmbg-path', trust_remote_code=True, local_files_only=True) + pool.add.assert_called_once_with('background', model) + assert pool.sealed is True + assert resident.config == {'rmbg_model': 'verified-rmbg-path'} + + +def test_foreground_controller_is_shared_and_never_prepares_native_or_full_models(monkeypatch): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setattr(release_backend, '_foreground_started', False) + monkeypatch.setattr(release_backend, '_foreground_worker', None) + monkeypatch.setattr(release_backend, '_foreground_snapshot', {'stage': 'pending'}) + updates = Mock(wraps=release_backend._update_foreground) + monkeypatch.setattr(release_backend, '_update_foreground', updates) + stage = Mock(return_value=Path('verified-rmbg-path')) + preload = Mock(return_value={'loaded_models': 1}) + monkeypatch.setattr(release_backend, '_stage_foreground_models', stage) + monkeypatch.setattr(zerogpu_runtime, 'preload_foreground', preload) + monkeypatch.setattr(release_backend.native_build, 'start', Mock(side_effect=AssertionError('native startup'))) + monkeypatch.setattr(release_backend, 'start_preparation', Mock(side_effect=AssertionError('full startup'))) + with ThreadPoolExecutor(max_workers=2) as executor: + workers = list(executor.map(lambda _: release_backend.start_foreground_preparation(), range(2))) + assert workers[0] is workers[1] + workers[0].join(timeout=5) + assert not workers[0].is_alive() + stage.assert_called_once_with() + preload.assert_called_once() + assert release_backend.foreground_status()['stage'] == 'ready_released' + phases = [call.kwargs['startup_progress'] for call in updates.call_args_list + if 'startup_progress' in call.kwargs] + assert all(phase['total'] > 0 and 0 <= phase['completed'] <= phase['total'] for phase in phases) + assert [(phase['unit'], phase['completed'], phase['total']) for phase in phases] == [ + ('steps', 0, 1), ('models', 0, 1), ('models', 1, 1)] + + +def test_foreground_dispatch_succeeds_while_full_models_are_unprepared(monkeypatch, tmp_path): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setattr(release_backend, 'model_status', lambda: {'stage': 'pending'}) + monkeypatch.setattr(release_backend, 'foreground_status', lambda: {'stage': 'ready_released'}) + foreground = Mock(return_value={'loaded_models': 1}) + monkeypatch.setattr(zerogpu_runtime, 'foreground', foreground) + assert release_backend.run_resident(tmp_path, 0, 120, None, operation='remove_background') == {'loaded_models': 1} + foreground.assert_called_once_with(str(tmp_path)) + + +def test_foreground_and_later_full_preload_keep_separate_residents(monkeypatch, tmp_path): + from hf_space import runtime_worker + def candidate(count): + return SimpleNamespace(prepare=Mock(), pool=SimpleNamespace( + sealed=True, summary=lambda: {'loaded_models': count}, restore_all=Mock())) + cutout, full = candidate(1), candidate(18) + monkeypatch.setattr(zerogpu_runtime, '_resident', None) + monkeypatch.setattr(zerogpu_runtime, '_foreground_resident', None) + foreground, generate = Mock(), Mock() + monkeypatch.setattr(runtime_worker, 'execute_foreground', foreground) + monkeypatch.setattr(runtime_worker, 'execute', generate) + zerogpu_runtime.preload_foreground('rmbg', Mock(), factory=lambda _: cutout) + assert zerogpu_runtime._resident is None + assert zerogpu_runtime._run(str(tmp_path), 0, foreground=True) == {'loaded_models': 1} + foreground.assert_called_once_with(cutout, tmp_path) + zerogpu_runtime.preload('full-config', Mock(), factory=lambda _: full) + assert zerogpu_runtime._run(str(tmp_path), 42) == {'loaded_models': 18} + generate.assert_called_once_with(full, tmp_path, 42, foreground_prepared=False) + assert zerogpu_runtime._foreground_resident is cutout + zerogpu_runtime.clear() + assert zerogpu_runtime._foreground_resident is cutout + cutout.pool.restore_all.assert_called_once_with() + full.pool.restore_all.assert_called_once_with() diff --git a/tests/test_zerogpu_split_runtime.py b/tests/test_zerogpu_split_runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..fd5dba3e39ace5ac33cf9a5d02ad439fa941199c --- /dev/null +++ b/tests/test_zerogpu_split_runtime.py @@ -0,0 +1,329 @@ +"""Geometry is durable before the independently requested material GPU lease.""" +import importlib.util +import io +import json +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import numpy as np +from PIL import Image +import pytest + +from hf_space import release_backend, runtime_worker, split_pipeline, zerogpu_runtime + + +def frontend_modules(monkeypatch): + """Load real adapters without the package's unrelated eager Torch imports.""" + root = Path(__file__).resolve().parents[1] / 'supratex' + modules = {} + for name, path in (('supratex', root), ('supratex.frontend', root/'frontend'), + ('supratex.inference', root/'inference')): + module = ModuleType(name) + module.__path__ = [str(path)] + monkeypatch.setitem(sys.modules, name, module) + modules[name] = module + if '.' in name: + setattr(modules[name.rsplit('.', 1)[0]], name.rsplit('.', 1)[1], module) + for suffix in ('weights', 'frontend.camera', 'frontend.firered', 'frontend.preprocess', + 'frontend.geometry', 'frontend.qwen_image', 'frontend.pipeline', 'inference.runtime'): + name = 'supratex.' + suffix + spec = importlib.util.spec_from_file_location(name, root.joinpath(*suffix.split('.')).with_suffix('.py')) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + modules[name] = module + parent, leaf = name.rsplit('.', 1) + setattr(modules[parent], leaf, module) + modules['supratex.inference'].run_completion = Mock() + return (modules['supratex.frontend.pipeline'], modules['supratex.inference.runtime'], + modules['supratex.weights'], modules['supratex.inference']) + + +def checkpoint(folder, seed=42, depth=False): + folder.mkdir(parents=True) + for name in split_pipeline.CHECKPOINT_FILES[:-1]: + (folder / name).write_bytes(('immutable-' + name).encode()) + files = list(split_pipeline.CHECKPOINT_FILES[:-1]) + if depth: + np.save(folder/'image_depth.npy', np.array([[1, np.nan], [2, 3]], dtype=np.float32), allow_pickle=False) + files.append('image_depth.npy') + (folder / 'checkpoint.json').write_text(json.dumps({ + 'format': split_pipeline.CHECKPOINT_FORMAT, 'seed': seed, + 'remove_floaters_applied': True, 'files': files})) + + +@pytest.mark.parametrize('options, depth_expected', [({}, True), ({'export_image_depth': False}, False), + ({'record_material_provenance': True}, False)]) +def test_geometry_frontend_returns_before_any_material_inference(monkeypatch, tmp_path, options, depth_expected): + pipeline, runtime, weights, _inference = frontend_modules(monkeypatch) + monkeypatch.setattr(runtime, 'configure_inference_runtime', Mock()) + monkeypatch.setattr(weights, 'resolve_model_dir', lambda source, **_: Path(source)) + monkeypatch.setattr(pipeline, 'validate_qwen', lambda *a, **k: (tmp_path, tmp_path, tmp_path/'adapter', None)) + material = Mock(side_effect=AssertionError('Geometry must stop before material inference')) + monkeypatch.setattr(pipeline, 'predict_material_maps', material) + + def prepare(_input, *, output_dir, **options): + rgba, rgb = output_dir/'condition_rgba.png', output_dir/'condition_rgb.png' + Image.new('RGBA', (8, 8), (255, 255, 255, 255)).save(rgba) + Image.new('RGB', (8, 8), 'white').save(rgb) + return rgba, rgb, 'prepared' + + monkeypatch.setattr(pipeline, 'prepare_image', prepare) + def camera(*args, depth_output=None, **kwargs): + if depth_output is not None: + np.save(depth_output, np.ones((8, 8), dtype=np.float32), allow_pickle=False) + return {'distance': 2, 'camera_angle_x': .5} + estimate = Mock(side_effect=camera) + monkeypatch.setattr(pipeline, 'estimate', estimate) + + def geometry(_package, _model, _input, destination, **options): + destination.mkdir() + files = {'geometry': 'geometry.pt', 'camera': 'camera.json', 'receipt': 'geometry_receipt.json'} + for key, name in files.items(): + (destination/name).write_text('optimized-camera' if key == 'camera' else name) + return {key: str(destination/name) for key, name in files.items()} + + generate = Mock(side_effect=geometry) + monkeypatch.setattr(pipeline, 'run_geometry', generate) + result = pipeline.run(tmp_path/'input.png', tmp_path/'checkpoint', moge_model=tmp_path, + geometry_package=tmp_path, geometry_model=tmp_path, material_base=tmp_path, + qwen_runtime=tmp_path, resident=object(), foreground_prepared=True, geometry_only=True, **options) + estimate.assert_called_once() + generate.assert_called_once() + material.assert_not_called() + assert Path(result['camera']).read_text() == 'optimized-camera' + assert (tmp_path/'checkpoint'/'geometry_receipt.json').is_file() + assert not (tmp_path/'checkpoint'/'material').exists() + assert (tmp_path/'checkpoint'/'image_depth.npy').exists() is depth_expected + assert ('image_depth' in result) is depth_expected + assert (estimate.call_args.kwargs['depth_output'] is not None) is depth_expected + + +@pytest.mark.parametrize('depth', [False, True]) +def test_texture_consumes_exact_checkpoint_without_geometry_or_camera(monkeypatch, tmp_path, depth): + pipeline, runtime, _weights, inference = frontend_modules(monkeypatch) + checkpoint(tmp_path / split_pipeline.CHECKPOINT_DIRECTORY, depth=depth) + source = tmp_path / split_pipeline.CHECKPOINT_DIRECTORY + before = {path.name: path.read_bytes() for path in source.iterdir()} + for name in ('run', 'run_geometry', 'estimate', 'prepare_image'): + monkeypatch.setattr(pipeline, name, Mock(side_effect=AssertionError('Geometry was rerun'))) + monkeypatch.setattr(runtime, 'configure_inference_runtime', Mock()) + + def material(rgba, rgb, output, **options): + assert rgba.read_bytes() == before['condition_rgba.png'] + assert rgb.read_bytes() == before['condition_rgb.png'] + output.mkdir() + for name in ('albedo.png', 'mra.png'): + (output/name).write_bytes(b'material-map') + return output/'albedo.png', output/'mra.png' + + monkeypatch.setattr(pipeline, 'predict_material_maps', Mock(side_effect=material)) + completed = [] + def completion(**options): + handoff = options['input_dir'] + assert (handoff/'geometry.pt').read_bytes() == before['geometry.pt'] + assert (handoff/'camera.json').read_bytes() == before['camera.json'] + assert (handoff/'image_depth.npy').exists() is depth + if depth: + assert (handoff/'image_depth.npy').read_bytes() == before['image_depth.npy'] + assert options['remove_floaters'] is False + assert options['preview_only'] is True + assert (handoff/'albedo.png').read_bytes() == b'material-map' + completed.append(options) + monkeypatch.setattr(inference, 'run_completion', completion) + args = SimpleNamespace(output=tmp_path/'asset.npz', seed=42, material_base='base', + material_model='adapter', qwen_runtime=tmp_path, qwen_deps=None, material_revision=None, + cache_dir=None, model_config='config', checkpoint='weights', trellis_root='vae', + texture_size=4096, decimation_target=500000, texture_format='png', gutter=False) + result = split_pipeline.run_texture(args, resident=object()) + assert len(completed) == 1 and result['geometry_reused'] is True + assert {name: (source/name).read_bytes() for name in before} == before + for name in ('run', 'run_geometry', 'estimate', 'prepare_image'): + getattr(pipeline, name).assert_not_called() + + +def test_checkpoint_requires_complete_contract_and_original_seed(tmp_path): + folder = tmp_path/'checkpoint' + checkpoint(folder) + with pytest.raises(ValueError, match='different seed'): + split_pipeline.copy_checkpoint(folder, tmp_path/'copy', seed=43) + assert not (tmp_path/'copy').exists() + (folder/'camera.json').unlink() + with pytest.raises(ValueError, match='incomplete'): + split_pipeline.validate_checkpoint(folder) + + +def test_declared_depth_must_exist_and_undeclared_depth_is_not_silently_copied(tmp_path): + new = tmp_path/'new'; checkpoint(new, depth=True) + (new/'image_depth.npy').unlink() + with pytest.raises(ValueError, match='incomplete'): + split_pipeline.copy_checkpoint(new, tmp_path/'missing-copy') + assert not (tmp_path/'missing-copy').exists() + old = tmp_path/'old'; checkpoint(old) + np.save(old/'image_depth.npy', np.ones((2, 2), dtype=np.float32), allow_pickle=False) + with pytest.raises(ValueError, match='undeclared depth'): + split_pipeline.copy_checkpoint(old, tmp_path/'unbound-copy') + assert not (tmp_path/'unbound-copy').exists() + + +@pytest.mark.parametrize('options, enabled', [({}, True), ({'export_image_depth': False}, False), + ({'record_material_provenance': True}, False)]) +def test_full_frontend_depth_defaults_and_dual_view_exclusion(monkeypatch, tmp_path, options, enabled): + pipeline, runtime, weights, _ = frontend_modules(monkeypatch) + monkeypatch.setattr(runtime, 'configure_inference_runtime', Mock()) + monkeypatch.setattr(weights, 'resolve_model_dir', lambda source, **_: Path(source)) + def prepare(_input, *, output_dir, **kwargs): + rgba, rgb = output_dir/'condition_rgba.png', output_dir/'condition_rgb.png' + Image.new('RGBA', (2, 2), 'white').save(rgba) + Image.new('RGB', (2, 2), 'white').save(rgb) + return rgba, rgb, 'alpha' + def estimate(*args, depth_output=None, **kwargs): + if depth_output is not None: + np.save(depth_output, np.ones((2, 2), dtype=np.float32), allow_pickle=False) + return {'distance': 2, 'camera_angle_x': .5} + def material(rgba, rgb, output, **kwargs): + output.mkdir() + for name in ('albedo.png', 'mra.png'): + Image.new('RGB', (2, 2), 'white').save(output/name) + return output/'albedo.png', output/'mra.png' + monkeypatch.setattr(pipeline, 'prepare_image', prepare) + monkeypatch.setattr(pipeline, 'estimate', Mock(side_effect=estimate)) + monkeypatch.setattr(pipeline, 'predict_material_maps', Mock(side_effect=material)) + result = pipeline.run(tmp_path/'input.png', tmp_path/'full', moge_model=tmp_path, + material_backend='firered', material_base=tmp_path, **options) + pipeline.estimate.assert_called_once(); pipeline.predict_material_maps.assert_called_once() + assert (tmp_path/'full'/'image_depth.npy').exists() is enabled + assert ('image_depth' in result) is enabled + assert Path(result['albedo']).is_file() and Path(result['mra']).is_file() + + +@pytest.mark.parametrize('mode', ['default', 'legacy-args', '--no-filter', '--no-image-depth', 'dual-view']) +def test_infer_entrypoint_depth_flag_reaches_frontend_with_resident_compatibility(monkeypatch, tmp_path, mode): + pipeline, runtime, _weights, inference = frontend_modules(monkeypatch) + path = Path(__file__).resolve().parents[1]/'infer.py' + spec = importlib.util.spec_from_file_location('_test_infer_depth', path) + module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module) + monkeypatch.setattr(module, '_validate_inference_options', Mock()) + argv = ['--output', str(tmp_path/'asset.glb'), '--moge-model', str(tmp_path), + '--geometry-package', str(tmp_path), '--geometry-model', str(tmp_path), + '--pretrained', str(tmp_path), '--qwen-runtime', str(tmp_path)] + if mode.startswith('--'): + argv.append(mode) + args = module.parse_args(argv) + assert args.no_filter is mode.startswith('--') + if mode == 'legacy-args': + del args.no_filter + args.dual_view = mode == 'dual-view'; args.back_view_dir = tmp_path/'back' + monkeypatch.setattr(module, '_validate_input_output', lambda *_: (tmp_path/'input.png', args.output, tmp_path/'work', True)) + monkeypatch.setattr(module, '_completion_paths', lambda *_: (tmp_path/'config', tmp_path/'weights', tmp_path/'vae')) + monkeypatch.setattr(runtime, 'configure_inference_runtime', Mock(return_value={'low_vram': True})) + monkeypatch.setattr(sys.modules['supratex.frontend.geometry'], 'validate_source', Mock()) + monkeypatch.setattr(sys.modules['supratex.frontend.qwen_image'], 'validate', Mock()) + monkeypatch.setattr(pipeline, 'run', Mock(return_value={'camera': 'unchanged-camera'})) + monkeypatch.setattr(inference, 'run_completion', Mock(side_effect=lambda **kwargs: kwargs['output'].write_bytes(b'output'))) + resident = None if args.dual_view else SimpleNamespace(pool=SimpleNamespace(restore_all=Mock())) + module.run(args, resident=resident) + enabled = mode in ('default', 'legacy-args') + assert pipeline.run.call_args.kwargs['export_image_depth'] is enabled + assert pipeline.run.call_args.kwargs.get('record_material_provenance', False) is args.dual_view + if resident is not None: + assert pipeline.run.call_args.kwargs['resident'] is resident + resident.pool.restore_all.assert_called_once() + record = json.loads((tmp_path/'work'/'run.json').read_text()) + assert record['image_depth_filter'] == {'requested': not mode.startswith('--'), 'enabled': enabled, + 'disabled_reason': 'user_disabled' if mode.startswith('--') else 'dual_view_unsupported' if args.dual_view else None} + + +def test_geometry_and_texture_register_distinct_gpu_leases(monkeypatch, tmp_path): + registrations, entered = {}, [] + def gpu(*, duration): + def decorate(function): + registrations[function.__name__] = duration + def wrapped(*args, **kwargs): + entered.append(function.__name__) + return function(*args, **kwargs) + return wrapped + return decorate + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setitem(sys.modules, 'spaces', SimpleNamespace(GPU=gpu)) + spec = importlib.util.spec_from_file_location('_test_split_gpu_runtime', zerogpu_runtime.__file__) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + restore = Mock() + module._resident = SimpleNamespace(pool=SimpleNamespace(restore_all=restore, summary=lambda: {})) + geometry, texture = Mock(), Mock() + monkeypatch.setattr(runtime_worker, 'execute_geometry', geometry) + monkeypatch.setattr(runtime_worker, 'execute_texture', texture) + monkeypatch.setattr(runtime_worker, 'execute', Mock(side_effect=AssertionError('Legacy full generation'))) + module.geometry(str(tmp_path), 42, True) + texture.assert_not_called() + module.texture(str(tmp_path), 42) + assert entered == ['geometry', 'texture'] + assert registrations['geometry'] == 180 and registrations['texture'] == 180 + geometry.assert_called_once_with(module._resident, tmp_path, 42, foreground_prepared=True) + texture.assert_called_once_with(module._resident, tmp_path, 42) + assert restore.call_count == 2 + + +@pytest.mark.parametrize('stage', ['geometry', 'texture']) +def test_each_split_failure_restores_its_model_pool(monkeypatch, tmp_path, stage): + restore = Mock() + monkeypatch.setattr(zerogpu_runtime, '_resident', SimpleNamespace(pool=SimpleNamespace(restore_all=restore))) + monkeypatch.setattr(runtime_worker, 'execute_'+stage, Mock(side_effect=RuntimeError('stage failed'))) + with pytest.raises(RuntimeError, match='stage failed'): + zerogpu_runtime._run(str(tmp_path), 42, stage=stage) + restore.assert_called_once() + + +def test_backend_texture_copies_checkpoint_before_its_separate_lease(monkeypatch, tmp_path): + monkeypatch.setenv('SPACES_ZERO_GPU', 'true') + monkeypatch.setattr(release_backend, 'JOBS_ROOT', tmp_path/'jobs') + monkeypatch.setattr(release_backend, 'model_status', lambda: {'stage':'ready_released'}) + monkeypatch.setattr(release_backend.generation_progress, 'update', Mock()) + image = Image.new('RGBA', (8, 8), 'white') + calls = [] + geometry_bytes = {} + def lease(job, seed, timeout, token, *, operation, **options): + calls.append(operation) + geometry = operation == 'generate_geometry' + if geometry: + checkpoint(job/split_pipeline.CHECKPOINT_DIRECTORY, seed) + geometry_bytes.update({name:(job/split_pipeline.CHECKPOINT_DIRECTORY/name).read_bytes() + for name in split_pipeline.CHECKPOINT_FILES}) + np.savez(job/'geometry.npz', vertices=np.zeros((3,3)), faces=np.zeros((1,3), dtype=int)) + else: + assert operation == 'generate_texture' + for name, contents in geometry_bytes.items(): + assert (job/split_pipeline.CHECKPOINT_DIRECTORY/name).read_bytes() == contents + np.savez(job/'asset.npz', vertices=np.zeros((3,3)), faces=np.zeros((1,3), dtype=int), + coords=np.zeros((1,3)), attrs=np.zeros((1,6)), resolution=1024, origin=np.zeros(3), voxel_size=1/1024) + (job/'views').mkdir() + image.convert('RGB').save(job/'views'/'clay-00.webp') + (job/'preview.json').write_text(json.dumps({'views':1, 'modes':['clay'], 'default_mode':'clay', 'automatic_glb_export':False})) + (job/'stage-result.json').write_text(json.dumps({'status':'complete', 'output_kind':'geometry' if geometry else 'final_ovoxel'})) + return {'loaded_models':18} + monkeypatch.setattr(release_backend, 'run_resident', lease) + source, _, report = release_backend.generate_geometry(image, 42, foreground_prepared=True) + assert calls == ['generate_geometry'] and report['output_kind'] == 'geometry' + result, _, report = release_backend.generate_texture(Path(source).parent, 42) + assert calls == ['generate_geometry', 'generate_texture'] + assert Path(result).name == 'asset.npz' and report['geometry_reused'] is True + + +def test_dedicated_worker_supports_split_operations_without_legacy_generate(monkeypatch, tmp_path): + monkeypatch.setattr(runtime_worker, 'JOBS_ROOT', tmp_path) + jobs = [tmp_path/('a'*32), tmp_path/('b'*32)] + for job in jobs: + job.mkdir() + pool = SimpleNamespace(sealed=True, summary=lambda:{'loaded_models':18}, restore_all=Mock()) + resident = SimpleNamespace(pool=pool, prepare=Mock()) + geometry, texture, full = Mock(), Mock(), Mock(side_effect=AssertionError('Full generation called')) + stream = io.StringIO(''.join(json.dumps({'job':str(job), 'epoch':'test', 'seed':42, 'operation':operation})+'\n' + for job, operation in zip(jobs, ('generate_geometry','generate_texture')))) + assert runtime_worker.serve('config', tmp_path/'status.json', 'test', stream, + factory=lambda _:resident, run_job=full, run_geometry=geometry, run_texture=texture) == 0 + geometry.assert_called_once() + texture.assert_called_once() + assert pool.restore_all.call_count == 2