icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add cutout shell and volume node assets for depth volume processing

Browse Source 
- 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:
马南
2026-09-08 17:18:03 +08:00
parent 7dec3b7ca4
commit 196b964efc
71 changed files with 3440 additions and 1988 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
src/anyimage/assets/O_AnyImage.blend LFS
View File

Binary file not shown.

6
src/anyimage/common/node.py
View File

@@ -10,10 +10,12 @@ def node_asset_path():
return Path(__file__).resolve().parents[1] / "assets" / NODE_ASSET_NAME
def load_node_group(group_name):
def load_node_group(group_name, *, asset_revision=None):
existing = bpy.data.node_groups.get(group_name)
if existing is not None:
return existing
if asset_revision is None or existing.get("anyimage_asset_revision") == asset_revision:
return existing
existing.name = group_name + " (Saved)"
asset_path = node_asset_path()
if not asset_path.is_file():

4
src/anyimage/common/object.py
View File

@@ -3,11 +3,11 @@
import bpy
def modifier_input_identifier(node_group, name):
def modifier_input_identifier(node_group, name, *, subtype=None):
for item in node_group.interface.items_tree:
if item.item_type != "SOCKET" or item.in_out != "INPUT":
continue
if item.name == name:
if item.name == name and (subtype is None or getattr(item, "subtype", None) == subtype):
return item.identifier
raise RuntimeError(f"Geometry Nodes input is missing: {name}")

5
src/anyimage/operators/cutout_tool/fine_mesh.py
View File

@@ -563,6 +563,11 @@ def triangulate(contours, spine, path_edges, free, spacing):
edges.extend(((boundary_count + a, boundary_count + b) for a, b in path_edges))
result = delaunay_2d_cdt([Vector(v) for v in p], edges, [], 0, 1e-07, True)
points = np.asarray([tuple(v) for v in result[0]])
# Restore input precision before testing area and containment. Float32 CDT
# coordinates can turn collinear outline samples into artificial slivers.
for vertex, originals in enumerate(result[3]):
if originals:
points[vertex] = p[originals].mean(axis=0)
faces = np.asarray(result[2], dtype=np.int32)
faces = faces[g.points_inside_contours(points[faces].mean(axis=1), contours)]
area = np.abs(

40
src/anyimage/operators/cutout_tool/geometry.py
View File

@@ -1238,6 +1238,35 @@ class BaseShape:
balloon: object
def refine_balloon_triangles(points, faces, boundary):
"""Add interior samples to boundary ears and interior boundary chords."""
from collections import Counter
import numpy as np
edges = Counter(tuple(sorted((face[i], face[(i + 1) % 3]))) for face in faces for i in range(3))
values, midpoints = list(points), {}
for edge, count in edges.items():
if count == 2 and all(boundary[index] for index in edge):
midpoints[edge] = len(values)
values.append(points[list(edge)].mean(axis=0))
refined = []
for face in faces:
ring = []
for index, vertex in enumerate(face):
ring.append(vertex)
midpoint = midpoints.get(tuple(sorted((vertex, face[(index + 1) % 3]))))
if midpoint is not None:
ring.append(midpoint)
if len(ring) == 3 and not all(boundary[vertex] for vertex in face):
refined.append(tuple(face))
continue
center = len(values)
values.append(points[list(face)].mean(axis=0))
refined.extend((center, ring[i], ring[(i + 1) % len(ring)]) for i in range(len(ring)))
points = np.asarray(values)
return points, refined, boundary_vertices_from_faces(len(points), refined)
def cutout_domain(
selection_contours,
bounds,
@@ -1276,6 +1305,17 @@ def cutout_domain(
(start, end, faces.tolist(), pixel_to_planar(contour))
for start, end, faces, contour in pixel_components
]
refined_points, refined_faces, refined_components = [], [], []
offset = 0
for start, end, component_faces, contour in components:
points, faces, _ = refine_balloon_triangles(planar[start:end], component_faces, boundary[start:end])
refined_points.append(points)
refined_faces.extend(tuple(offset + index for index in face) for face in faces)
refined_components.append((offset, offset + len(points), faces, contour))
offset += len(points)
planar = np.concatenate(refined_points)
planar_faces, components = refined_faces, refined_components
boundary = boundary_vertices_from_faces(len(planar), planar_faces)
profile = balloon_profile(
balloon_algorithm, planar, planar_faces, boundary, components
)

6
src/anyimage/operators/cutout_tool/interaction.py
View File

@@ -4,10 +4,10 @@ from ...properties import ai_ready
from .shape import CUTOUT_SHAPE_LABELS
CUTOUT_SHAPE_BUTTONS = (
("SURFACE", "MESH_GRID"),
("BALLOON", "MESH_UVSPHERE"),
("FLAT", "MESH_GRID"),
("SOLID", "MESH_UVSPHERE"),
("DEPTH_BALLOON", "META_BALL"),
("DEPTH_SURFACE", "SURFACE_NCURVE"),
("DEPTH_SOLID", "SURFACE_NCURVE"),
)
CUTOUT_MESH_DETAIL_SPACING = {

10
src/anyimage/operators/cutout_tool/main.py
View File

@@ -73,15 +73,17 @@ _HIDDEN = {"HIDDEN", "SKIP_SAVE"}
def normal_mode_for_shape(shape, generate_normal):
if shape == "DEPTH_SOLID":
return NormalMode.OBJECT
if shape == "DEPTH_BALLOON":
return NormalMode.TANGENT
if not generate_normal:
return NormalMode.NONE
if shape == "DEPTH_SURFACE":
return NormalMode.OBJECT
return NormalMode.TANGENT
def depth_mode_for_shape(shape):
if shape == "DEPTH_SURFACE":
if shape == "DEPTH_SOLID":
return "VECTOR"
if shape == "DEPTH_BALLOON":
return "Z"
@@ -200,7 +202,7 @@ class CutoutSelectionToShape(JobOperatorBase, bpy.types.Operator):
source_object_name: bpy.props.StringProperty(options=_HIDDEN)
selection_path_json: bpy.props.StringProperty(options=_HIDDEN)
selection_bounds_json: bpy.props.StringProperty(options=_HIDDEN)
shape: bpy.props.StringProperty(options=_HIDDEN)
shape: bpy.props.StringProperty(default="SOLID", options=_HIDDEN)
point_spacing: bpy.props.FloatProperty(
default=32.0,
min=0.01,

27
src/anyimage/operators/cutout_tool/object.py
View File

@@ -31,10 +31,9 @@ from .geometry import (
)
from .shape import (
BALLOON_ATTRIBUTE_NAME,
BALLOON_CUTOUT_SHAPES,
NORMAL_SCALE_ATTRIBUTE_NAME,
CUTOUT_X_TO_Z,
CUTOUT_NODE_GROUP_NAME,
CUTOUT_SHAPE_INDICES,
CUTOUT_NODE_GROUP_NAMES,
DEPTH_CUTOUT_SHAPES,
)
@@ -131,11 +130,11 @@ def create_cropped_color_image(source_object, crop_bounds, source_rgba):
return image
def load_cutout_node_group():
return load_node_group(CUTOUT_NODE_GROUP_NAME)
def load_cutout_node_group(shape):
return load_node_group(CUTOUT_NODE_GROUP_NAMES[shape], asset_revision=11)
def _connect_balloon_normal_scale(material):
def _connect_cutout_normal_scale(material):
nodes = material.node_tree.nodes
image_layer = next(
node
@@ -146,7 +145,7 @@ def _connect_balloon_normal_scale(material):
)
attribute = nodes.new("ShaderNodeAttribute")
attribute.attribute_type = "GEOMETRY"
attribute.attribute_name = BALLOON_ATTRIBUTE_NAME
attribute.attribute_name = NORMAL_SCALE_ATTRIBUTE_NAME
attribute.location = (-550.0, -250.0)
material.node_tree.links.new(
attribute.outputs["Fac"], image_layer.inputs["Normal Scale"]
@@ -206,8 +205,7 @@ def create_shape_object(
scene=context.scene,
)
depth_reference = None
if shape in BALLOON_CUTOUT_SHAPES:
_connect_balloon_normal_scale(material)
_connect_cutout_normal_scale(material)
uniform_scale = None
depth_direction = None
if shape in DEPTH_CUTOUT_SHAPES:
@@ -227,7 +225,7 @@ def create_shape_object(
depth_reference = DEFAULT_DEPTH_BASELINE
if shape == "DEPTH_BALLOON":
depth_direction = DEFAULT_DEPTH_DIRECTION
elif shape == "DEPTH_SURFACE":
elif shape == "DEPTH_SOLID":
depth_reference = metric_depth_reference(
camera_depth,
calibration_valid,
@@ -268,22 +266,19 @@ def create_shape_object(
inward_axis,
)
mesh_object["anyimage_mesh_shape"] = shape
node_group = load_cutout_node_group()
node_group = load_cutout_node_group(shape)
modifier = mesh_object.modifiers.new(name=node_group.name, type="NODES")
modifier.node_group = node_group
modifier.show_group_selector = False
for name, value in (
("Shape", CUTOUT_SHAPE_INDICES[shape]),
(
"Inward Axis",
0 if inward_axis == CutoutInwardAxis.NEGATIVE_Z else 1,
),
("Cleanup", cleanup_value),
("Smooth Iterations", 0),
("Smooth", 1.0),
("Cleanup Threshold", cleanup_value),
(
"Thickness",
1.0 if shape in {"BALLOON", "DEPTH_BALLOON"} else 0.0,
0.0 if shape == "FLAT" else 1.0,
),
):
set_modifier_input(

20
src/anyimage/operators/cutout_tool/shape.py
View File

@@ -1,7 +1,7 @@
"""Cutout Shape identifiers shared by runtime and node-asset tools."""
BALLOON_ATTRIBUTE_NAME = "o_balloon"
CUTOUT_NODE_GROUP_NAME = "O Image Cutout"
NORMAL_SCALE_ATTRIBUTE_NAME = "o_normal_scale"
CUTOUT_X_TO_Z = (
(0.0, -1.0, 0.0, 0.0),
(0.0, 0.0, 1.0, 0.0),
@@ -15,20 +15,16 @@ CUTOUT_Z_TO_X = (
(0.0, 0.0, 0.0, 1.0),
)
CUTOUT_NODE_GROUP_NAMES = {
"SURFACE": "O Image Surface",
"BALLOON": "O Image Balloon",
"FLAT": "O Image Cutout",
"SOLID": "O Image Cutout",
"DEPTH_BALLOON": "O Image Depth Balloon",
"DEPTH_SURFACE": "O Image Depth Surface",
"DEPTH_SOLID": "O Image Depth Cutout",
}
CUTOUT_SHAPE_LABELS = {
"SURFACE": "Surface",
"BALLOON": "Balloon",
"FLAT": "Flat",
"SOLID": "Solid",
"DEPTH_BALLOON": "Depth Balloon",
"DEPTH_SURFACE": "Depth Surface",
}
CUTOUT_SHAPE_INDICES = {
shape: index for index, shape in enumerate(CUTOUT_NODE_GROUP_NAMES)
"DEPTH_SOLID": "Depth Solid",
}
CUTOUT_SHAPES = frozenset(CUTOUT_NODE_GROUP_NAMES)
DEPTH_CUTOUT_SHAPES = frozenset({"DEPTH_BALLOON", "DEPTH_SURFACE"})
BALLOON_CUTOUT_SHAPES = frozenset({"BALLOON", "DEPTH_BALLOON"})
DEPTH_CUTOUT_SHAPES = frozenset({"DEPTH_BALLOON", "DEPTH_SOLID"})

4
src/anyimage/operators/upscale.py
View File

@@ -29,7 +29,7 @@ from ..common.image import (
_HIDDEN = {"HIDDEN", "SKIP_SAVE"}
UPSCALE_FACTOR = 4
UPSCALE_FACTOR = 2
def upscale_output_size(image):
@@ -49,7 +49,7 @@ def upscale_input_size_allowed(image, max_input_size):
class UpscaleImage(JobOperatorBase, bpy.types.Operator):
bl_idname = "anyimage.upscale_image"
bl_label = "Upscale"
bl_description = "Upscale this Image Empty by 4x with an AI model"
bl_description = "Upscale this Image Empty by 2x with an AI model"
bl_options = {"UNDO"}
job_type = "upscale-image"
starting_message = "Starting Image Empty upscale"

4
src/anyimage/properties.py
View File

@@ -251,8 +251,8 @@ def cutout_generate_normal_property():
return bpy.props.BoolProperty(
name="Normal Map",
description=(
"Generate a MoGe-2 normal texture; Depth Surface uses Object Space "
"and other shapes use Tangent Space"
"Generate a normal map for Flat and Solid; "
"Depth Solid and Depth Balloon always generate normal maps"
),
default=False,
)

31
src/anyimage/server/model_catalog.py
View File

@@ -99,7 +99,7 @@ REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec(
directory_name="realesr-general-wdn-x4v3",
repository="xinntao/Real-ESRGAN",
license_name="BSD 3-Clause",
description="Fast and conservative 4x image upscaler using ONNX Runtime",
description="Fast and conservative 2x output from a 4x ONNX model",
supports_gs=False,
download_patterns=("realesr-general-wdn-x4v3.onnx",),
ready_patterns=("realesr-general-wdn-x4v3.onnx",),
@@ -117,7 +117,7 @@ SAFMN_MODEL = ModelSpec(
directory_name="safmn-real-x4",
repository="sunny2109/SAFMN",
license_name="Apache 2.0",
description="Balanced real-world 4x image upscaler using ONNX Runtime",
description="Balanced 2x output from a real-world 4x ONNX model",
supports_gs=False,
download_patterns=("safmn-real-x4.onnx",),
ready_patterns=("safmn-real-x4.onnx",),
@@ -136,7 +136,7 @@ REALESRGAN_X4PLUS = ModelSpec(
repository="jonathanst29/tinier-upscale-models",
license_name="BSD 3-Clause",
huggingface_repository="jonathanst29/tinier-upscale-models",
description="Detailed general-purpose 4x image upscaler using ONNX Runtime",
description="Detailed 2x output from a general-purpose 4x ONNX model",
supports_gs=False,
download_patterns=("realesrgan-x4plus.onnx",),
ready_patterns=("realesrgan-x4plus.onnx",),
@@ -149,12 +149,37 @@ REALESRGAN_X4PLUS = ModelSpec(
)
PISA_SR = ModelSpec(
key="PISA_SR",
label="PiSA-SR",
directory_name="pisa-sr",
repository="csslc/PiSA-SR",
license_name="Apache 2.0 / CreativeML Open RAIL++-M",
description="Generative 2x upscale with fixed detail strength 0.7",
supports_gs=False,
download_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
ready_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
r2_directory="pisa-sr",
r2_files=(
("encoder.onnx", 68426712, "4b93e4a3d81c3349a78ab96619501e39ee8f815d14a20b73fc4f969e30a1ceeb"),
("decoder.onnx", 99084589, "c9021384c402e267654907527fa5f0ee86cab1814bfd96233216a15ae8e83c4c"),
("unet-pixel.onnx", 1732713007, "4f381574ad584050d0534c8cca794fb5828cb8e99e148f8efc1aaac20833c0b1"),
("unet-full.onnx", 1732713007, "3b0fd44ff1093608a65865c3f26d0060775337e1fc117d8ac29aacbe748d3f9c"),
("noise.npy", 32896, "a94ed01e2e963022cfa8ec01e04d2e96e599408ef0c385878347eaacd25a663b"),
("LICENSE-APACHE.txt", 11358, "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"),
("LICENSE-MODEL.txt", 14206, "e514164be7f2ca2c0fb39e7283b01ab4a06d802837eaf171aa4645cfd7ee1157"),
("NOTICE.txt", 1081, "ccb55fcb98cce42401368b88e3ccd555f4ff49ac1347c6ce063061dd48726050"),
),
)
UPSCALE_MODELS = {
model.key: model
for model in (
REALESRGAN_GENERAL_WDN_X4V3,
SAFMN_MODEL,
REALESRGAN_X4PLUS,
PISA_SR,
)
}

32
src/anyimage/server/model_manager.py
View File

@@ -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):

153
src/anyimage/server/models/onnx_pisa_sr.py Normal file
View File

@@ -0,0 +1,153 @@
"""Fixed-strength PiSA-SR inference with one resident ONNX component."""
import gc
from pathlib import Path
from tempfile import TemporaryDirectory
import numpy as np
from PIL import Image
from .onnx_runtime import create_session, run_session
from .onnx_upscale import RESOURCE_ERROR_TOKENS, UpscaleResourceError
TILE_SIZE = 512
OVERLAP = 128
LATENT_SCALE = np.float16(0.18215)
SEMANTIC_STRENGTH = np.float16(0.7)
def tile_starts(length):
padded = max(TILE_SIZE, (length + 7) // 8 * 8)
starts = list(range(0, padded - TILE_SIZE + 1, TILE_SIZE - OVERLAP))
if starts[-1] != padded - TILE_SIZE:
starts.append(padded - TILE_SIZE)
return starts, padded
def correct_color(pixels, source):
target = pixels.astype(np.float32) / 255
reference = np.asarray(source, dtype=np.float32) / 255
axes = (0, 1)
deviation = np.sqrt(target.var(axis=axes, ddof=1) + 1e-5)
reference_deviation = np.sqrt(reference.var(axis=axes, ddof=1) + 1e-5)
corrected = (target - target.mean(axis=axes)) / deviation
corrected = corrected * reference_deviation + reference.mean(axis=axes)
return Image.fromarray((corrected.clip(0, 1) * 255).astype(np.uint8))
class PisaSession:
def __init__(self, directory, device):
if device == "coreml":
raise RuntimeError("PiSA-SR CoreML is not validated; select CPU for this model")
self.directory = Path(directory)
self.device = device
self.session = None
self.component = None
self.noise = np.load(self.directory / "noise.npy", allow_pickle=False)
if self.noise.shape != (1, 4, 64, 64) or not np.isfinite(self.noise).all():
raise ValueError("Invalid PiSA-SR noise asset")
def close(self):
self.session = None
self.component = None
gc.collect()
def activate(self, component):
if self.component == component:
return
self.close()
self.session = create_session(self.directory / f"{component}.onnx", self.device)
self.session.disable_fallback()
if "CoreMLExecutionProvider" in self.session.get_providers():
self.close()
raise RuntimeError("PiSA-SR CoreML is not validated; select CPU for this model")
self.component = component
def run(self, value, final=False):
try:
result = run_session(
self.session, ["output"],
{"input": np.ascontiguousarray(value, dtype=np.float16)},
release_memory=final,
)[0]
except Exception as error:
if "CUDAExecutionProvider" in self.session.get_providers():
raise RuntimeError(
"PiSA-SR failed on CUDA; select DirectML or CPU. " + str(error)
) from error
raise
if not np.isfinite(result).all():
raise RuntimeError("PiSA-SR returned non-finite pixels on the selected device")
return result
def infer(model, image, cancel_check=None, *, release_memory=True):
def check():
if cancel_check is not None and cancel_check():
raise RuntimeError("Task cancelled")
source = image.convert("RGB").resize(
(image.width * 2, image.height * 2), Image.Resampling.BICUBIC,
)
rows, height = tile_starts(source.height)
columns, width = tile_starts(source.width)
coordinates = [(y, x) for y in rows for x in columns]
array = np.asarray(source)
array = np.pad(array, ((0, height-source.height), (0, width-source.width), (0, 0)),
mode="reflect" if min(source.size) > 1 else "edge")
try:
with TemporaryDirectory(prefix="anyimage-pisa-") as temporary:
directory = Path(temporary)
check()
model.activate("encoder")
for index, (y, x) in enumerate(coordinates):
check()
tile = array[y:y+TILE_SIZE, x:x+TILE_SIZE].astype(np.float32) / 255
tensor = (tile.transpose(2, 0, 1)[None] * 2 - 1).astype(np.float16)
mean, logvar = np.split(model.run(tensor, index == len(coordinates)-1), 2, axis=1)
noise = np.take(np.take(model.noise, (np.arange(64)+y//8) % 64, axis=2),
(np.arange(64)+x//8) % 64, axis=3)
std = np.exp(logvar.clip(-30, 20) * np.float16(0.5)).astype(np.float16)
latent = ((mean + std * noise) * LATENT_SCALE).astype(np.float16)
np.save(directory / f"latent-{index}.npy", latent)
for component in ("unet-pixel", "unet-full"):
check()
model.activate(component)
for index in range(len(coordinates)):
check()
latent = np.load(directory / f"latent-{index}.npy")
prediction = model.run(latent, index == len(coordinates)-1)
pixel_path = directory / f"pixel-{index}.npy"
if component == "unet-pixel":
np.save(pixel_path, prediction)
else:
pixel = np.load(pixel_path)
corrected = latent - (pixel + SEMANTIC_STRENGTH * (prediction-pixel))
np.save(directory / f"latent-{index}.npy", corrected)
pixel_path.unlink()
check()
model.activate("decoder")
result = np.zeros((height, width, 3), dtype=np.float32)
weights = np.zeros((height, width, 1), dtype=np.float32)
ramp = np.minimum(np.arange(TILE_SIZE)+1, np.arange(TILE_SIZE, 0, -1))
ramp = np.minimum(ramp / OVERLAP, 1).astype(np.float32)
blend = (ramp[:, None] * ramp[None, :])[..., None]
for index, (y, x) in enumerate(coordinates):
check()
latent = np.load(directory / f"latent-{index}.npy")
decoded = model.run(latent, release_memory and index == len(coordinates)-1)
pixels = decoded[0].transpose(1, 2, 0).astype(np.float32)
result[y:y+TILE_SIZE, x:x+TILE_SIZE] += pixels * blend
weights[y:y+TILE_SIZE, x:x+TILE_SIZE] += blend
check()
result = (result / weights)[:source.height, :source.width]
pixels = ((result.clip(-1, 1)*0.5+0.5)*255).clip(0, 255).astype(np.uint8)
return correct_color(pixels, source)
except (MemoryError, UnicodeDecodeError) as error:
model.close()
raise UpscaleResourceError("PiSA-SR exhausted available inference memory") from error
except Exception as error:
model.close()
if any(token in str(error).lower() for token in RESOURCE_ERROR_TOKENS):
raise UpscaleResourceError("PiSA-SR exhausted available inference memory") from error
raise

7
src/anyimage/server/models/onnx_runtime.py
View File

@@ -62,9 +62,16 @@ def create_session(
elif provider == "CoreMLExecutionProvider" and coreml_options:
provider = (provider, dict(coreml_options))
configured_providers.append(provider)
session_options = {}
if "DmlExecutionProvider" in providers:
options = onnxruntime.SessionOptions()
options.enable_mem_pattern = False
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
session_options["sess_options"] = options
return onnxruntime.InferenceSession(
str(path),
providers=configured_providers,
**session_options,
)

7
src/anyimage/server/models/upscale.py
View File

@@ -19,6 +19,9 @@ def infer_one(
from .onnx_upscale import UpscaleResourceError, infer
if parameters["model"] == "PISA_SR":
from .onnx_pisa_sr import infer
model_directory = model_manager.directory(parameters["model"])
model = model_manager.get_upscale(
model_directory,
@@ -37,6 +40,10 @@ def infer_one(
cancel_check=cancel_check,
release_memory=release_memory,
)
prediction = prediction.resize(
(source.width * 2, source.height * 2),
Image.Resampling.LANCZOS,
)
if source_alpha is None:
return prediction
alpha = source_alpha.resize(