| """Reuse curated examples or validated inference results before requesting a GPU.""" |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import shutil |
| import time |
| import uuid |
|
|
| from hf_space import example_assets |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| EXAMPLE_IMAGE = ROOT/'assets/pitcher_plant.png' |
| CACHE = Path('/data/generated/pbr-results-v1') |
| EXAMPLE_NAMES = example_assets.EXAMPLE_NAMES |
|
|
|
|
| def revision(): |
| return os.environ.get('SUPRATEX_MODEL_REVISION', '499db28a11b614873a2cfbcfd1990024ba0f1d34') |
|
|
|
|
| def image_digest(image): |
| rgba = image.convert('RGBA') |
| digest = hashlib.sha256(f'RGBA:{rgba.width}:{rgba.height}:'.encode()) |
| digest.update(rgba.tobytes()) |
| return digest.hexdigest() |
|
|
|
|
| def cache_key(image, seed): |
| request = {'input_rgba_sha256': image_digest(image), 'seed': int(seed), |
| 'project_revision': revision(), 'profile': 'torch28-sdpa-pbr4096-v1', |
| 'project_repository': os.environ.get('SUPRATEX_MODEL_REPO', 'supra3d/supra-tex')} |
| return hashlib.sha256(json.dumps(request, sort_keys=True).encode()).hexdigest(), request |
|
|
|
|
| def _lookup_saved(key, request): |
| try: |
| record = json.loads((CACHE/(key+'.json')).read_text()) |
| model = CACHE/(key+'.glb') |
| if record.get('request') == request and example_assets.valid_model(model, record['sha256'], record['glb_bytes']): |
| return str(model) |
| except (OSError, ValueError, KeyError): |
| pass |
| return None |
|
|
|
|
| def default_preview(): |
| return example_assets.lookup_name('pitcher_plant') |
|
|
|
|
| def lookup(image, seed): |
| if image is None: |
| return None |
| name = example_assets.identify(image, seed) |
| if name: |
| |
| return example_assets.lookup_name(name) |
| key, request = cache_key(image, seed) |
| return _lookup_saved(key, request) |
|
|
|
|
| def prepare_examples(): |
| example_assets.prepare() |
|
|
|
|
| def examples_status(): |
| return example_assets.status() |
|
|
|
|
| def store(image, seed, source, report): |
| key, request = cache_key(image, seed) |
| CACHE.mkdir(parents=True, exist_ok=True) |
| model = CACHE/(key+'.glb') |
| temporary = CACHE/(key+'.'+uuid.uuid4().hex+'.part') |
| try: |
| shutil.copyfile(source, temporary) |
| if not example_assets.valid_model(temporary, report['sha256'], report['glb_bytes']): |
| raise ValueError('PBR result cache checksum mismatch') |
| temporary.replace(model) |
| finally: |
| temporary.unlink(missing_ok=True) |
| record = {'request': request, 'sha256': report['sha256'], 'glb_bytes': report['glb_bytes'], |
| 'source_job': report.get('job_directory'), 'created_unix': time.time()} |
| temporary = CACHE/(key+'.'+uuid.uuid4().hex+'.json.part') |
| try: |
| temporary.write_text(json.dumps(record, indent=2)) |
| temporary.replace(CACHE/(key+'.json')) |
| finally: |
| temporary.unlink(missing_ok=True) |
| return str(model) |
|
|
