| """Persist final o-voxels and their rendered views in the mounted data Bucket."""
|
| import json
|
| import hashlib
|
| import logging
|
| import math
|
| from pathlib import Path
|
| import re
|
| import shutil
|
| import struct
|
| import tempfile
|
| import threading
|
| import time
|
| import uuid
|
| from html import escape
|
| from urllib.parse import quote
|
| from concurrent.futures import ThreadPoolExecutor
|
|
|
| PREVIEW_MODES = (
|
| ('normal', 'Normal', 'normal.png'),
|
| ('clay', 'Clay', 'clay.png'),
|
| ('base_color', 'Base color', 'basecolor.png'),
|
| ('shaded_studio', 'HDRI Studio', 'hdri_studio.png'),
|
| ('shaded_forest', 'HDRI Forest', 'hdri_forest.png'),
|
| ('shaded_sunset', 'HDRI Sunset', 'hdri_sunset.png'),
|
| ('shaded_courtyard', 'HDRI Courtyard', 'hdri_courtyard.png'),
|
| ('shaded', 'Studio', 'hdri_studio.png'),
|
| )
|
|
|
| ROOT = Path('/data/generated/ovoxel-results-v1')
|
| JOBS = Path('/tmp/supratex-jobs')
|
| PREVIEW_CACHE = JOBS/'preview-cache'
|
| PREVIEW_CACHE_BYTES = 512*1024*1024
|
| PREVIEW_CACHE_ENTRIES = 64
|
| 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')
|
| 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 '
|
| 'material maps were not saved. Chart-isolated texture gutters were applied.'
|
| )
|
|
|
|
|
| def directory(identifier):
|
| if not isinstance(identifier, str) or re.fullmatch(r'[0-9a-f]{32}', identifier) is None:
|
| raise ValueError('Invalid generated asset identifier')
|
| return ROOT/identifier
|
|
|
|
|
| def local_manifest(folder):
|
| """Describe immutable display assets while their bytes are still local."""
|
| folder = Path(folder)
|
| version = hashlib.sha256(str(folder).encode()).hexdigest()
|
| metadata = json.loads((folder/'preview.json').read_text())
|
| files = {f'frames/{path.name}': {'size':path.stat().st_size, 'version':version}
|
| for path in (folder/'views').glob('*.webp')}
|
| return files, metadata
|
|
|
|
|
| def persist(source, report):
|
| source = Path(source)
|
| identifier = source.parent.name
|
| destination = directory(identifier)
|
| destination.mkdir(parents=True, exist_ok=True)
|
| started = time.monotonic()
|
| for name in ('asset.npz','input.png','report.json','preview.json','timings.jsonl'):
|
| shutil.copyfile(source.parent/name, destination/name)
|
| if (source.parent/'input-camera.json').is_file():
|
| shutil.copyfile(source.parent/'input-camera.json', destination/'input-camera.json')
|
| 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'
|
| views.mkdir(exist_ok=True)
|
| frames = list((source.parent/'views').glob('*.webp'))
|
| def copy_frame(path):
|
| shutil.copyfile(path, views/path.name)
|
|
|
| with ThreadPoolExecutor(max_workers=2, thread_name_prefix='preview-publication') as pool:
|
| list(pool.map(copy_frame, frames))
|
| files, metadata = local_manifest(source.parent)
|
| (destination/'display-manifest.json').write_text(json.dumps(
|
| {'files':files,'preview_metadata':metadata,'save_seconds':time.monotonic()-started}))
|
| (destination/'ready.json').write_text(json.dumps({'complete':True,'identifier':identifier,
|
| 'project_revision':report['project_revision'],'output_kind':'final_ovoxel',
|
| 'reference_inputs_saved':reference_saved}))
|
| return identifier
|
|
|
|
|
| def _file_digest(path):
|
| digest = hashlib.sha256()
|
| with path.open('rb') as stream:
|
| while chunk := stream.read(1024*1024):
|
| digest.update(chunk)
|
| 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
|
| 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
|
| names = (*REFERENCE_FILES, *(name for name in REFERENCE_OPTIONAL_FILES if (folder/name).is_file()))
|
| for name in names:
|
| path = folder/name
|
| if (path.is_symlink() or any(parent.is_symlink() for parent in path.absolute().parents)
|
| or not path.is_file() or path.stat().st_size <= 0):
|
| raise ValueError('The saved reference inputs are incomplete or invalid')
|
| camera = json.loads((folder/'camera.json').read_text(encoding='utf-8'))
|
| try:
|
| fov, distance, scale = (float(camera['camera_angle_x']), float(camera['distance']),
|
| float(camera.get('mesh_scale', 1.0)))
|
| except (KeyError, TypeError, ValueError) as error:
|
| raise ValueError('The saved reference camera is invalid') from error
|
| if not (all(math.isfinite(value) for value in (fov, distance, scale))
|
| and 0 < fov < math.pi and distance > 0 and scale > 0):
|
| raise ValueError('The saved reference camera is invalid')
|
| for name in ('condition_rgba.png', 'albedo.png', 'mra.png',
|
| *(('condition_rgb.png',) if 'condition_rgb.png' in names else ())):
|
| with Image.open(folder/name) as image:
|
| size = image.size
|
| if ((name in ('albedo.png', 'mra.png') and size != (1024, 1024))
|
| or (name.startswith('condition_') and (size[0] != size[1] or size[0] <= 0))
|
| or (name == 'condition_rgba.png' and image.mode != 'RGBA')):
|
| raise ValueError('The saved reference image dimensions are invalid')
|
| image.verify()
|
| if 'condition_rgb.png' in names:
|
| 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')
|
| return actual
|
|
|
|
|
| def _persist_reference_inputs(source, destination):
|
| bindings = _validate_reference_inputs(source)
|
| final = destination/REFERENCE_DIRECTORY
|
| if final.exists():
|
| _validate_reference_inputs(final, bindings)
|
| return
|
| staged = destination/('.reference-input.'+uuid.uuid4().hex+'.partial')
|
| staged.mkdir()
|
| try:
|
| for name in bindings:
|
| shutil.copyfile(source/name, staged/name)
|
| _validate_reference_inputs(staged, bindings)
|
| (staged/'manifest.json').write_text(json.dumps({'version':1, 'files':bindings}, indent=2))
|
| staged.replace(final)
|
| finally:
|
| if staged.exists():
|
| if (staged.resolve().parent != destination.resolve()
|
| or re.fullmatch(r'\.reference-input\.[0-9a-f]{32}\.partial', staged.name) is None):
|
| raise ValueError('Invalid temporary reference input path')
|
| shutil.rmtree(staged)
|
|
|
|
|
| def _reference_inputs(identifier, root, work, token):
|
| """Recover durable observations, or a still-local pre-upgrade generation."""
|
| durable = root/REFERENCE_DIRECTORY
|
| local = JOBS/identifier/'asset_work'/'input'
|
| if durable.is_dir():
|
| manifest = json.loads((durable/'manifest.json').read_text(encoding='utf-8'))
|
| bindings = manifest.get('files', {})
|
| if (manifest.get('version') != 1 or not isinstance(bindings, dict)
|
| or not set(REFERENCE_FILES) <= set(bindings)
|
| or not set(bindings) <= set(REFERENCE_FILES+REFERENCE_OPTIONAL_FILES)
|
| or any(not isinstance(value, str) or re.fullmatch(r'[0-9a-f]{64}', value) is None
|
| for value in bindings.values())):
|
| raise ValueError('The saved reference input manifest is invalid')
|
| if local.is_dir():
|
| try:
|
| _validate_reference_inputs(local, bindings)
|
| return local
|
| except (OSError, ValueError):
|
| pass
|
| staged = work/REFERENCE_DIRECTORY
|
| staged.mkdir()
|
| for name in bindings:
|
| _copy_progress(durable/name, staged/name, token)
|
| _validate_reference_inputs(staged, bindings)
|
| return staged
|
| if local.is_dir() and all((local/name).is_file() for name in REFERENCE_FILES):
|
|
|
|
|
| _persist_reference_inputs(local, root)
|
| return local
|
| if json.loads((root/'ready.json').read_text(encoding='utf-8')).get('reference_inputs_saved') is True:
|
| raise ValueError('The saved reference inputs are unavailable. Please retry when storage is available.')
|
| return None
|
|
|
|
|
| def _frame_names(metadata):
|
| modes, views = metadata.get('modes'), metadata.get('views')
|
| known_modes = {mode for mode, _, _ in PREVIEW_MODES}
|
| if (not isinstance(modes, list) or not modes
|
| or any(not isinstance(mode, str) for mode in modes)
|
| or len(modes) != len(set(modes))
|
| or any(mode not in known_modes for mode in modes)
|
| or not isinstance(views, int) or not 1 <= views <= 64):
|
| raise ValueError('Invalid saved preview metadata')
|
| return [f'{mode}-{view:02d}.webp' for mode in modes for view in range(views)]
|
|
|
|
|
| def preview_view_order(metadata):
|
| """Keep the slider starting at the input direction for saved previews too."""
|
| _frame_names(metadata)
|
| count = metadata['views']
|
|
|
|
|
| offset = 4 if count == 16 and not metadata.get('view_layout') else 0
|
| return [(view + offset) % count for view in range(count)]
|
|
|
|
|
| def _complete_preview(root, metadata, names):
|
| try:
|
| if json.loads((root/'preview.json').read_text()) != metadata:
|
| return False
|
| return all((root/'views'/name).is_file() and
|
| (root/'views'/name).stat().st_size > 0 for name in names)
|
| except (OSError, ValueError):
|
| return False
|
|
|
|
|
| def _remove_cached_preview(path):
|
|
|
| if path.resolve().parent != PREVIEW_CACHE.resolve():
|
| raise ValueError('Invalid preview cache path')
|
| shutil.rmtree(path)
|
|
|
|
|
| def _trim_preview_cache(keep):
|
| entries = []
|
| for child in PREVIEW_CACHE.iterdir():
|
| if (child.is_dir() and child.resolve().parent == PREVIEW_CACHE.resolve()
|
| and re.fullmatch(r'[0-9a-f]{32}-[0-9a-f]{16}', child.name)):
|
| size = sum(path.stat().st_size for path in child.rglob('*') if path.is_file())
|
| entries.append((child.stat().st_mtime, child, size))
|
| total = sum(size for _, _, size in entries)
|
| count = len(entries)
|
| for _, child, size in sorted(entries):
|
| if total <= PREVIEW_CACHE_BYTES and count <= PREVIEW_CACHE_ENTRIES:
|
| break
|
| if child != keep:
|
| _remove_cached_preview(child)
|
| total -= size
|
| count -= 1
|
|
|
|
|
| def _preview_root(identifier):
|
| """Never give Gradio file requests a path on the network-mounted bucket."""
|
| root = directory(identifier)
|
| metadata = json.loads((root/'preview.json').read_text())
|
| names = _frame_names(metadata)
|
| local = JOBS/identifier
|
| if _complete_preview(local, metadata, names):
|
| return local, metadata
|
| payload = json.dumps(metadata, sort_keys=True, separators=(',', ':')).encode()
|
| cached = PREVIEW_CACHE/f'{identifier}-{hashlib.sha256(payload).hexdigest()[:16]}'
|
| with _preview_lock:
|
| if _complete_preview(cached, metadata, names):
|
| cached.touch()
|
| return cached, metadata
|
| if not (root/'ready.json').is_file():
|
| raise ValueError('The generated preview has not finished saving')
|
| PREVIEW_CACHE.mkdir(parents=True, exist_ok=True)
|
| staged = Path(tempfile.mkdtemp(prefix=f'.{identifier}-', dir=PREVIEW_CACHE))
|
| try:
|
| (staged/'views').mkdir()
|
| total = 0
|
| for name in names:
|
| destination = staged/'views'/name
|
| shutil.copyfile(root/'views'/name, destination)
|
| size = destination.stat().st_size
|
| total += size
|
| if size == 0 or total > MAX_PREVIEW_BYTES:
|
| raise ValueError('Saved preview is empty or exceeds its cache limit')
|
| (staged/'preview.json').write_bytes(payload)
|
| if json.loads((root/'preview.json').read_text()) != metadata:
|
| raise ValueError('Saved preview changed while loading')
|
| if cached.exists():
|
| _remove_cached_preview(cached)
|
| staged.replace(cached)
|
| _trim_preview_cache(cached)
|
| return cached, metadata
|
| finally:
|
| if staged.exists():
|
| _remove_cached_preview(staged)
|
|
|
|
|
| def preview_html(identifier):
|
| root, metadata = _preview_root(identifier)
|
| default = metadata.get('default_mode', 'shaded' if 'shaded' in metadata['modes'] else 'shaded_forest')
|
| if default not in metadata['modes']:
|
| default = metadata['modes'][0]
|
| images = []
|
| for mode in metadata['modes']:
|
| for view, source_view in enumerate(preview_view_order(metadata)):
|
| url = '/gradio_api/file='+quote(str(root/'views'/f'{mode}-{source_view:02d}.webp'),safe='/')
|
| active = mode == default and view == 0
|
| source = f' src="{url}"' if active else ''
|
| hidden = '' if active else ' hidden'
|
| images.append(f'<img{source} data-src="{url}" decoding="async" data-mode="{mode}" data-view="{view}" alt="{mode.replace("_"," ")} view {view+1}"{hidden}>')
|
| buttons = []
|
| for mode, label, icon in PREVIEW_MODES:
|
| if mode not in metadata['modes']:
|
| continue
|
|
|
| import base64
|
| icon_path = Path(__file__).resolve().parents[1]/'assets/preview-lighting'/icon
|
| swatch = base64.b64encode(icon_path.read_bytes()).decode('ascii')
|
| active = mode == default
|
| buttons.append(f'<button type="button" data-mode="{mode}" class="{"active" if active else ""}" '
|
| f'title="{label}" aria-label="{label}" aria-pressed="{str(active).lower()}">'
|
| f'<img src="data:image/png;base64,{swatch}" alt=""><span class="st-mode-name">{label}</span></button>')
|
| return (f'<div class="st-ovoxel-preview" data-mode="{escape(default, quote=True)}" data-view="0">'
|
| '<div class="st-ovoxel-frames">'+''.join(images)+
|
| '<span class="st-ovoxel-status" role="status" aria-live="polite" hidden '
|
| 'style="position:absolute;bottom:8px;left:8px;right:8px;padding:8px;'
|
| 'font-size:12px;color:#ddd;background:#222d;border-radius:6px"></span></div>'
|
| '<div class="st-ovoxel-modes" role="group" aria-label="Render mode">'+''.join(buttons)+'</div>'
|
| f'<label class="st-ovoxel-angle">View angle <input type="range" min="0" max="{metadata["views"]-1}" value="0" step="1" aria-label="View angle"></label></div>')
|
|
|
|
|
| def extract(identifier, extraction_token=None, *, prepare_download=None):
|
| from hf_space import extraction_progress as progress
|
| from hf_space.timing import observe_stages
|
| token = progress.start(extraction_token)
|
| stages = []
|
| started = time.monotonic()
|
| def stage(name):
|
| if name in progress.STAGES:
|
| stages.append({'stage': name, 'start_seconds': round(time.monotonic()-started, 3)})
|
| progress.stage(token, name)
|
| try:
|
| with observe_stages(stage):
|
| output = _extract(identifier, token, stage=stage, stages=stages,
|
| prepare_download=prepare_download)
|
| progress.finish(token)
|
| return output
|
| except Exception:
|
| progress.finish(token, failed=True)
|
| raise
|
|
|
|
|
| def _copy_progress(source, destination, token):
|
| """Use one sequential transfer instead of ZIP-member reads on the Bucket."""
|
| from hf_space import extraction_progress as progress
|
| total = source.stat().st_size
|
| completed, notified = 0, time.monotonic()
|
| progress.transfer(token, 0, total)
|
| with source.open('rb') as reader, destination.open('xb') as writer:
|
| while chunk := reader.read(8*1024*1024):
|
| writer.write(chunk)
|
| completed += len(chunk)
|
| now = time.monotonic()
|
| if now-notified >= .25 or completed == total:
|
| progress.transfer(token, completed, total)
|
| notified = now
|
| if completed != total:
|
| raise RuntimeError('The saved 3D artifact changed during transfer')
|
|
|
|
|
| def _complete_glb(path):
|
| try:
|
| with path.open('rb') as stream:
|
| magic, version, length = struct.unpack('<4sII', stream.read(12))
|
| return magic == b'glTF' and version == 2 and length == path.stat().st_size
|
| except (OSError, ValueError, struct.error):
|
| return False
|
|
|
|
|
| def _local_asset(identifier, remote):
|
| """Reuse the immutable file we just generated, without another Bucket read."""
|
| local = JOBS/identifier/'asset.npz'
|
| if (not local.is_symlink() and not any(parent.is_symlink() for parent in local.absolute().parents)
|
| and local.is_file() and local.stat().st_size == remote.stat().st_size):
|
| return local
|
| return None
|
|
|
|
|
| def _extract(identifier, token, *, stage, stages, prepare_download=None):
|
| """Keep random NPZ reads, mesh writes and validation on local storage."""
|
| from hf_space.timing import observe_stages
|
| root = directory(identifier)
|
| if not (root/'ready.json').is_file():
|
| raise ValueError('The generated o-voxel has not finished saving')
|
| output = root/'asset.glb'
|
| if _complete_glb(output):
|
| if prepare_download:
|
| stage('download_cache')
|
| _prepare_download(prepare_download, output, output)
|
| return str(output)
|
| stage('load')
|
| started = time.monotonic()
|
| JOBS.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
| with tempfile.TemporaryDirectory(prefix='.glb-export-', dir=JOBS) as temporary:
|
| work = Path(temporary)
|
| source = _local_asset(identifier, root/'asset.npz')
|
| if source is None:
|
| source = work/'asset.npz'
|
| _copy_progress(root/'asset.npz', source, token)
|
| reference = _reference_inputs(identifier, root, work, token)
|
| local_output = work/'asset.glb'
|
|
|
|
|
| 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:
|
| postprocess = {'reprojection':'applied', 'gutter':'observation_guided',
|
| 'postprocess_warning':None}
|
| else:
|
| postprocess = {'reprojection':'skipped_missing_inputs', 'gutter':'chart_isolated',
|
| 'postprocess_warning':LEGACY_REPROJECTION_WARNING}
|
| stage('validate')
|
| from hf_space.backend import validate_glb
|
| metadata = validate_glb(local_output)
|
| stage('publish')
|
| staged = root/('.asset.'+uuid.uuid4().hex+'.glb.partial')
|
| try:
|
| _copy_progress(local_output, staged, token)
|
| staged.replace(output)
|
| finally:
|
| staged.unlink(missing_ok=True)
|
| if prepare_download:
|
| stage('download_cache')
|
| _prepare_download(prepare_download, local_output, output)
|
|
|
|
|
| report = {'status':'complete', 'device':'cuda',
|
| 'elapsed_seconds':time.monotonic()-started, 'stage_events':stages,
|
| 'texture_size':4096, 'target_faces':500000, **metadata, **postprocess}
|
| report_path = root/('.extraction.'+uuid.uuid4().hex+'.json.partial')
|
| try:
|
| report_path.write_text(json.dumps(report, indent=2))
|
| report_path.replace(root/'extraction.json')
|
| finally:
|
| report_path.unlink(missing_ok=True)
|
| return str(output)
|
|
|
|
|
| def _prepare_download(callback, local, durable):
|
| try:
|
| callback(local, durable)
|
| except Exception:
|
|
|
|
|
| LOG.warning('GLB was saved, but its download cache could not be prewarmed', exc_info=True)
|
|
|
|
|
| def _export_local(source, output, *, gutter=False):
|
| import gc
|
| import numpy as np
|
| import torch
|
| from supratex.representations import MeshWithVoxel
|
| from supratex.inference.pipeline import _export, LAYOUT, _runtime
|
| if not torch.cuda.is_available():
|
| raise RuntimeError('GLB extraction requires CUDA')
|
| with np.load(source,allow_pickle=False) as state:
|
| resolution = int(state['resolution'])
|
| representation = MeshWithVoxel(torch.from_numpy(state['vertices']).cuda(),
|
| torch.from_numpy(state['faces']).cuda(),state['origin'].tolist(),float(state['voxel_size']),
|
| torch.from_numpy(state['coords']).cuda(),torch.from_numpy(state['attrs']).cuda(),
|
| torch.Size([1,6,resolution,resolution,resolution]),layout=LAYOUT)
|
| try:
|
| _export(output, representation,4096,500000,'png',_runtime(),
|
| gutter=gutter, low_vram=True)
|
| finally:
|
|
|
|
|
| del representation
|
| gc.collect()
|
| torch.cuda.empty_cache()
|
|
|
