File size: 2,935 Bytes
a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff a60cf00 411b0ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """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:
# Never fall back to a superseded generated example.
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)
|