增加解码节点
This commit is contained in:
engineer6980
+2
-6
@@ -1,17 +1,13 @@
|
||||
# 导入你写的节点类
|
||||
from .cipher_nodes import ImageCipherEncode, ImageCipherDecode
|
||||
|
||||
# 建立 ComfyUI 内部识别名与 Python 类的映射
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageCipherEncode": ImageCipherEncode,
|
||||
"ImageCipherDecode": ImageCipherDecode
|
||||
}
|
||||
|
||||
# 建立节点在 UI 界面显示的中文/英文名称
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageCipherEncode": "🔒 RGB 行交织加密",
|
||||
"ImageCipherDecode": "🔓 RGB 行交织解密"
|
||||
"ImageCipherEncode": "🔒 Cipher Encode (Color Noise)",
|
||||
"ImageCipherDecode": "🔓 Cipher Decode (Preview)"
|
||||
}
|
||||
|
||||
# 必须导出这两个变量
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
+26
-37
@@ -1,7 +1,7 @@
|
||||
import torch
|
||||
|
||||
# ==========================================
|
||||
# 🔑 固定密钥：前端和PyCharm必须使用同一个数字
|
||||
# 🔑 全局密钥：必须保持一致
|
||||
# ==========================================
|
||||
SECRET_KEY = 999
|
||||
|
||||
@@ -14,68 +14,59 @@ class ImageCipherEncode:
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
# 输入标准 VAE 解码后的图像
|
||||
"images": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("cipher_color_noise",)
|
||||
RETURN_NAMES = ("cipher_image",)
|
||||
FUNCTION = "encode"
|
||||
CATEGORY = "CipherTools"
|
||||
|
||||
def encode(self, images):
|
||||
# images shape: [Batch, Height, Width, Channel=3]
|
||||
batch, h, w, c = images.shape
|
||||
out_list = []
|
||||
|
||||
# 使用固定密钥作为基础偏移量
|
||||
base_shift = SECRET_KEY
|
||||
|
||||
for i in range(batch):
|
||||
img = images[i] # [H, W, 3]
|
||||
img = images[i]
|
||||
|
||||
# 1. 拆分三个通道，它们现在的形状都是 [H, W]
|
||||
r_plane = img[:, :, 0]
|
||||
g_plane = img[:, :, 1]
|
||||
b_plane = img[:, :, 2]
|
||||
# 分离通道
|
||||
r = img[:, :, 0]
|
||||
g = img[:, :, 1]
|
||||
b = img[:, :, 2]
|
||||
|
||||
# 准备加密容器
|
||||
r_encrypted = torch.zeros_like(r_plane)
|
||||
g_encrypted = torch.zeros_like(g_plane)
|
||||
b_encrypted = torch.zeros_like(b_plane)
|
||||
# 准备容器
|
||||
r_enc = torch.zeros_like(r)
|
||||
g_enc = torch.zeros_like(g)
|
||||
b_enc = torch.zeros_like(b)
|
||||
|
||||
# 2. 核心加密循环：对每一行进行不同程度的位移
|
||||
# 加密循环：正向位移 (shifts=正数)
|
||||
for y in range(h):
|
||||
# 关键点：R, G, B 使用不同的倍率 (1, 2, 3)，确保它们错开
|
||||
# 计算 R 通道当前行的位移量
|
||||
shift_r = (y * base_shift * 1) % w
|
||||
r_encrypted[y, :] = torch.roll(r_plane[y, :], shifts=shift_r, dims=0)
|
||||
|
||||
# 计算 G 通道当前行的位移量 (倍率不同)
|
||||
shift_g = (y * base_shift * 2) % w
|
||||
g_encrypted[y, :] = torch.roll(g_plane[y, :], shifts=shift_g, dims=0)
|
||||
|
||||
# 计算 B 通道当前行的位移量 (倍率不同)
|
||||
shift_b = (y * base_shift * 3) % w
|
||||
b_encrypted[y, :] = torch.roll(b_plane[y, :], shifts=shift_b, dims=0)
|
||||
|
||||
# 3. 重新组合成彩色图像 [H, W, 3]
|
||||
# 此时 R,G,B 已经完全错位，图像看起来是彩色噪点
|
||||
encrypted_img = torch.stack((r_encrypted, g_encrypted, b_encrypted), dim=-1)
|
||||
out_list.append(encrypted_img)
|
||||
r_enc[y, :] = torch.roll(r[y, :], shifts=shift_r, dims=0)
|
||||
g_enc[y, :] = torch.roll(g[y, :], shifts=shift_g, dims=0)
|
||||
b_enc[y, :] = torch.roll(b[y, :], shifts=shift_b, dims=0)
|
||||
|
||||
out_list.append(torch.stack((r_enc, g_enc, b_enc), dim=-1))
|
||||
|
||||
return (torch.stack(out_list, dim=0),)
|
||||
|
||||
|
||||
# 解密节点（仅供 ComfyUI 内部测试用，实际业务在前端完成）
|
||||
class ImageCipherDecode:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"images": ("IMAGE",), }}
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("restored_image",)
|
||||
@@ -89,7 +80,7 @@ class ImageCipherDecode:
|
||||
|
||||
for i in range(batch):
|
||||
img = images[i]
|
||||
# 1. 拆分加密后的通道
|
||||
|
||||
r_enc = img[:, :, 0]
|
||||
g_enc = img[:, :, 1]
|
||||
b_enc = img[:, :, 2]
|
||||
@@ -98,19 +89,17 @@ class ImageCipherDecode:
|
||||
g_dec = torch.zeros_like(g_enc)
|
||||
b_dec = torch.zeros_like(b_enc)
|
||||
|
||||
# 2. 反向操作：向相反方向回滚
|
||||
# 解密循环：反向位移 (shifts=负数)
|
||||
for y in range(h):
|
||||
# 计算加密时用的位移量
|
||||
shift_r = (y * base_shift * 1) % w
|
||||
shift_g = (y * base_shift * 2) % w
|
||||
shift_b = (y * base_shift * 3) % w
|
||||
|
||||
# 使用负数 shifts 进行反向滚动
|
||||
# 注意这里的负号 -shift
|
||||
r_dec[y, :] = torch.roll(r_enc[y, :], shifts=-shift_r, dims=0)
|
||||
g_dec[y, :] = torch.roll(g_enc[y, :], shifts=-shift_g, dims=0)
|
||||
b_dec[y, :] = torch.roll(b_enc[y, :], shifts=-shift_b, dims=0)
|
||||
|
||||
restored_img = torch.stack((r_dec, g_dec, b_dec), dim=-1)
|
||||
out_list.append(restored_img)
|
||||
out_list.append(torch.stack((r_dec, g_dec, b_dec), dim=-1))
|
||||
|
||||
return (torch.stack(out_list, dim=0),)
|
||||
Reference in New Issue
Block a user