147 lines
5.9 KiB
Python
147 lines
5.9 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
from PIL import Image
|
|
|
|
# 获取当前文件所在目录，用于定位模型
|
|
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
MODELS_DIR = os.path.join(CURRENT_DIR, "onnx_models")
|
|
|
|
class PixelArtConverter:
|
|
def __init__(self):
|
|
self.sessions = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # ComfyUI 输入格式 [Batch, H, W, C]
|
|
"resolution": (["128", "256"], {"default": "128"}),
|
|
"temperature": ("FLOAT", {"default": 0.01, "min": 0.0001, "max": 1.0, "step": 0.01}),
|
|
"noise_amp": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "PixelArt"
|
|
|
|
def _get_session(self, model_key):
|
|
model_filename = f"pixelart_{model_key}.onnx"
|
|
model_path = os.path.join(MODELS_DIR, model_filename)
|
|
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"ONNX Model not found at: {model_path}")
|
|
|
|
if model_key not in self.sessions:
|
|
print(f"Loading ONNX model: {model_path}...")
|
|
# 优先尝试 CUDA，如果没有则 CPU
|
|
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
|
|
self.sessions[model_key] = ort.InferenceSession(model_path, providers=providers)
|
|
|
|
return self.sessions[model_key]
|
|
|
|
def process(self, image, resolution, temperature, noise_amp, seed):
|
|
# 1. 准备输入
|
|
# ComfyUI 图片是 Tensor [Batch, H, W, C]，范围 0-1
|
|
# 我们这里处理 Batch 中的每一张图
|
|
|
|
output_images = []
|
|
|
|
for i in range(image.shape[0]):
|
|
img_tensor = image[i] # [H, W, C]
|
|
|
|
# Tensor -> Numpy -> PIL (为了保持和原项目一致的 Resize 逻辑)
|
|
img_np = (img_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
|
img_pil = Image.fromarray(img_np)
|
|
|
|
# 处理 Alpha (如果是 RGBA 转为白底 RGB)
|
|
if img_pil.mode == 'RGBA':
|
|
bg = Image.new('RGBA', img_pil.size, (255, 255, 255, 255))
|
|
combined = Image.alpha_composite(bg, img_pil)
|
|
img_pil = combined.convert('RGB')
|
|
else:
|
|
img_pil = img_pil.convert('RGB')
|
|
|
|
# 2. 预处理 (Pre-processing)
|
|
w, h = img_pil.size
|
|
INPUT_SIZE = 1024
|
|
|
|
scale = INPUT_SIZE / max(h, w)
|
|
new_w = int(round(w * scale))
|
|
new_h = int(round(h * scale))
|
|
|
|
# Resize (LANCZOS)
|
|
img_resized = img_pil.resize((new_w, new_h), Image.Resampling.LANCZOS)
|
|
|
|
# 填充到 1024x1024 白底
|
|
canvas = Image.new("RGB", (INPUT_SIZE, INPUT_SIZE), (255, 255, 255))
|
|
pad_left = (INPUT_SIZE - new_w) // 2
|
|
pad_top = (INPUT_SIZE - new_h) // 2
|
|
canvas.paste(img_resized, (pad_left, pad_top))
|
|
|
|
# 归一化 (-1 到 1)
|
|
img_input_np = np.array(canvas).astype(np.float32)
|
|
img_input_np = (img_input_np / 255.0) * 2.0 - 1.0
|
|
# [H, W, C] -> [1, C, H, W]
|
|
img_input = img_input_np.transpose(2, 0, 1)[None, :, :, :]
|
|
|
|
# 3. 推理 (Inference)
|
|
sess = self._get_session(resolution)
|
|
|
|
# 随机种子处理
|
|
rng = np.random.default_rng(seed)
|
|
noise_input = rng.standard_normal((1, 32, 256)).astype(np.float32) * noise_amp
|
|
temp_input = np.array([temperature], dtype=np.float32)
|
|
|
|
input_feed = {
|
|
'input_image': img_input,
|
|
'input_noise': noise_input,
|
|
'temperature': temp_input
|
|
}
|
|
|
|
outputs = sess.run(None, input_feed)
|
|
pred_img = outputs[0][0] # [3, H, W]
|
|
|
|
# 4. 后处理 (Post-processing)
|
|
pred_img = pred_img.transpose(1, 2, 0) # [H, W, 3]
|
|
pred_img = np.clip(pred_img, 0.0, 1.0)
|
|
|
|
# 5. 裁剪逻辑 (Crop back)
|
|
# 计算裁剪区域
|
|
out_h_total, out_w_total = pred_img.shape[:2]
|
|
ratio = out_h_total / float(INPUT_SIZE)
|
|
|
|
crop_top = int(pad_top * ratio)
|
|
crop_left = int(pad_left * ratio)
|
|
crop_h = int(new_h * ratio)
|
|
crop_w = int(new_w * ratio)
|
|
|
|
# 确保至少有 1px
|
|
crop_h = max(1, crop_h)
|
|
crop_w = max(1, crop_w)
|
|
|
|
final_img_np = pred_img[crop_top : crop_top+crop_h, crop_left : crop_left+crop_w]
|
|
|
|
# 转回 Tensor [H, W, C]
|
|
output_tensor = torch.from_numpy(final_img_np)
|
|
output_images.append(output_tensor)
|
|
|
|
# 堆叠 Batch: [Batch, H, W, C]
|
|
# 注意：如果 Batch 内图片裁剪后尺寸不一致，torch.stack 会报错。
|
|
# 通常 ComfyUI 单次处理通常尺寸一致，这里为了稳健，如果尺寸不一致可以在这里做 Padding，
|
|
# 但最简单的做法是假设用户 Batch 输入尺寸长宽比一致。
|
|
if len(output_images) > 1:
|
|
try:
|
|
result = torch.stack(output_images, dim=0)
|
|
except RuntimeError:
|
|
# 如果尺寸不一致，无法堆叠，只能抛出第一张或者报错 (这是ComfyUI Tensor的限制)
|
|
print("Warning: Batch images have different aspect ratios after pixel art conversion. Returning first image only.")
|
|
result = output_images[0].unsqueeze(0)
|
|
else:
|
|
result = output_images[0].unsqueeze(0)
|
|
|
|
return (result,) |