Add cutout shell and volume node assets for depth volume processing
- Implemented `cutout_shell.py` to build front normals and shell positions for balloon geometry. - Created `cutout_volume.py` to establish smooth center fields and manage balloon surface attributes. - Introduced `pisa_sr_export.py` for exporting PiSA-SR components with dependencies and asset finalization.
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
import gc
|
||||
import hashlib
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
@@ -69,6 +70,9 @@ class ModelSessionCache:
|
||||
def get_upscale(self, model_dir, requested_device):
|
||||
from .models.onnx_upscale import create_session
|
||||
|
||||
if Path(model_dir).name == "pisa-sr":
|
||||
from .models.onnx_pisa_sr import PisaSession as create_session
|
||||
|
||||
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
|
||||
if self.upscale_model is not None and self.upscale_key == key:
|
||||
return self.upscale_model
|
||||
@@ -91,11 +95,17 @@ class ModelSessionCache:
|
||||
self._release_runtime_cache()
|
||||
|
||||
def release_upscale(self):
|
||||
close = getattr(self.upscale_model, "close", None)
|
||||
if close is not None:
|
||||
close()
|
||||
self.upscale_model = None
|
||||
self.upscale_key = None
|
||||
self._release_runtime_cache()
|
||||
|
||||
def close(self):
|
||||
close = getattr(self.upscale_model, "close", None)
|
||||
if close is not None:
|
||||
close()
|
||||
self.ben2_model = None
|
||||
self.da3_model = None
|
||||
self.moge_model = None
|
||||
@@ -124,6 +134,7 @@ class ModelManager:
|
||||
def __init__(self, models_directory):
|
||||
self.models_directory = Path(models_directory)
|
||||
self.sessions = ModelSessionCache()
|
||||
self.verified_files = {}
|
||||
|
||||
def directory(self, model_key):
|
||||
model = model_catalog.get_downloadable_model(model_key)
|
||||
@@ -132,6 +143,27 @@ class ModelManager:
|
||||
def ready(self, model_key):
|
||||
model = model_catalog.get_downloadable_model(model_key)
|
||||
directory = self.directory(model_key)
|
||||
if model_key == "PISA_SR":
|
||||
for filename, size, checksum in model.r2_files:
|
||||
path = directory / filename
|
||||
try:
|
||||
stat = path.stat()
|
||||
if stat.st_size != size:
|
||||
return False
|
||||
identity = (stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns)
|
||||
cached = self.verified_files.get(path)
|
||||
# Windows can coalesce timestamps for immediately repeated writes.
|
||||
cacheable = size >= 1024 * 1024 and time.time_ns() - stat.st_mtime_ns > 1_000_000_000
|
||||
if not cacheable or cached != (identity, checksum):
|
||||
with path.open("rb") as stream:
|
||||
actual = hashlib.file_digest(stream, "sha256").hexdigest()
|
||||
if actual != checksum:
|
||||
return False
|
||||
if cacheable:
|
||||
self.verified_files[path] = (identity, checksum)
|
||||
except OSError:
|
||||
return False
|
||||
return True
|
||||
return all(any(directory.glob(pattern)) for pattern in model.ready_patterns)
|
||||
|
||||
def get_ben2(self, model_dir, device):
|
||||
|
||||
Reference in New Issue
Block a user