nodes/Gue-e-ComfyUI-PanoCard
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

首次提交

Browse Source
This commit is contained in:
Gue
2025-03-18 14:19:20 +08:00
parent 8efe0e8b6f
commit c0bc1a0bb0
15 changed files with 59572 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+36
View File
@@ -0,0 +1,36 @@
from .pano import *
from .convert import *
__all__ = ['NODE_CLASS_MAPPINGS']
NODE_CLASS_MAPPINGS = {
"PanoCardViewer": PanoViewer,
"ImageWightPad": PanoImageWightPad,
"ImageHeightPad": PanoImageHeightPad,
"ImageCube2Equ": PanoImageCube2Equ,
"ImageEqu2Cube": PanoImageEqu2Cube,
"ImageEqu2Pic": PanoImageEqu2Pic,
"ImagePic2Equ": PanoImagePic2Equ,
"ImageEqu2Equ": PanoImageEqu2Equ,
"ImageClamp": PanoImageClamp,
"ImageOutClamp": PanoImageOutClamp,
"MaskOutClamp": PanoMaskOutClamp,
"PipePad": PanoImagePipe,
"ImagePad": PanoImagePad,
"ImageRoll": PanoImageRoll,
"MaskCondBatch":PanoDenseDiffCondBatch,
"ImageFaceCondBatch":PanoClipBatch,
"PromptSplit":PanoPromptSplit,
"RegionalPrompt":PanoRegionalPrompt,
"ImageEquSplit":PanoImageSplit,
"CondFaceSplit":PanoClipOutClamp,
"ImageAdjust":PanoImageAdjust,
"MaskFaceOutClamp": PanoMaskOutFaceClamp,
"ImageFaceToLong": PanoFaceToLong,
"CondFaceScheduleHookProvider": CondFaceScheduleHookProvider,
}
WEB_DIRECTORY = "./web"
__all__ = ['NODE_CLASS_MAPPINGS','NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
convert.py
+410
View File
@@ -0,0 +1,410 @@
import numpy as np
import equilib as py360
import torch
from .distort import plane_to_equirect
from .misc import *
import math
import nodes
def calculate_equ_heght(heght):
# 计算 sqrt(x * x * 3)
result = math.sqrt(heght * heght * 3)
# 将结果四舍五入到最接近的64的倍数
result = int((result / 64 + 0.5)) * 64
if result < 512:
result = 512
return result
def calculate_face_heght(heght):
# 计算 sqrt(x * x / 3)
result = math.sqrt(heght * heght / 3)
# 将结果四舍五入到最接近的64的倍数
result = int((result / 64 + 0.5)) * 64
if result < 512:
result = 512
return result
# Equirectangular全景图，Cubemap 立方体贴图 Perspective 透视图
class PanoImageCube2Equ:
def __init__(self):
pass
CATEGORY = "PanoCard/convert"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"height": ("INT", {
"default": 832,
"min": 0,
"step": 8
}),
"scale": ("FLOAT", {
"default": 0,
"max": 1.0,
"min": 0,
"step": 0.01
}),
"mask_split": ("BOOLEAN", {"default": True, "label_on": "true", "label_off": "false"}),
},
"optional": {
"face": ("IMAGE",),
"long": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE","MASK","MASK")
RETURN_NAMES = ("equ","mask","masks")
OUTPUT_IS_LIST = (False, False, True)
FUNCTION = "node"
def node(self, height, scale, mask_split, face = None, long = None):
sub_height = 0
result = np.array([])
order = "nearest"
if long is not None:
sub_height, sub_width = long.shape[1], long.shape[2]
if sub_height*6 != sub_width:
raise ValueError("传入图像宽高必须为6:1")
if height == 0:
height = calculate_equ_heght(sub_height)
long_face = long.permute(0, 3, 1, 2)
equ = py360.cube2equi(
cubemap = long_face,
height = height,
width = height * 2,
mode = order,
cube_format = 'horizon')
result = equ.permute(1, 2, 0).unsqueeze(0)
elif face is not None:
if not all(img.shape == face[0].shape for img in face):
raise ValueError("所有图像的尺寸必须相同")
sub_height = face[0].shape[0]
if height == 0:
height = calculate_equ_heght(sub_height)
face = tuple(face.permute(0, 3, 1, 2))
long_face = torch.cat(face, dim=2)
equ = py360.cube2equi(
cubemap = long_face,
height = height,
width = height * 2,
mode = order,
cube_format = 'horizon')
result = equ.permute(1, 2, 0).unsqueeze(0)
else:
if scale < 0.1:
raise ValueError("传入图像不能为空")
if height == 0:
height = 832
sub_height = height // 2
if scale > 0.01:
# 缩放并居中放置图像
def create_sub_image(height, scale, value):
H = int(height * scale)
# 计算目标位置
T = (height - H) // 2
# 创建目标大小的全零图像，形状为 [C, H, W]
result = torch.zeros((4, height, height), dtype=torch.uint8)
# 将缩放后的图像居中放置
result[:, T:T+H, T:T+H] = torch.full((4, H, H), value, dtype=torch.uint8)
return result
# 创建一个形状为 [C, H, W] 的组合图像
np_image = torch.zeros((4, sub_height, sub_height * 6), dtype=torch.uint8)
for i in range(6):
value = (i + 1) * 10
sub_image = create_sub_image(sub_height, scale, value)
# 将子图像拼接到 np_image 中
np_image[:, :, i * sub_height:(i + 1) * sub_height] = sub_image
equ = py360.cube2equi(
cubemap = np_image,
height = height,
width = height * 2,
mode = order,
cube_format='horizon')
equ_mask = equ.permute(1, 2, 0).unsqueeze(0)
# 转化输入为空时候，使用equ_mask返回
if np.size(result) == 0:
result = equ_mask / 255.0
result = result.to(torch.float32)
#提取透明通道转换为mask
mask = equ_mask[:, :, :, 3]
mask = (mask != 0).float()
#提取颜色通道转换为mask
masks = []
blue_channel = equ_mask[0][:, :, 2]
for i in range(6):
# 直接使用蓝色通道的值来创建遮罩
matching_pixels = (blue_channel == (i+1)*10).float()
# 转换为 PyTorch 张量
masks.append(matching_pixels)
masks = torch.stack(masks, dim=0)
if result.shape[-1] == 3:
result = add_alpha_channel(result)
if mask_split:
funmaskand = nodes.NODE_CLASS_MAPPINGS['BitwiseAndMask']()
mask = result[:, :, :, 3]
mask_a = [1.0 - mask] * 6
mask_a = torch.stack(mask_a, dim=0)
masks = funmaskand(mask_a, masks)
masks = masks.unsqueeze(0)
return (result, mask, masks)
else:
# 当scale为0时，返回全透明
mask = result[:, :, :, 1]
masks = torch.zeros(1, 6, 64, 64)
return (result, mask, masks)
class PanoImageEqu2Cube:
def __init__(self):
pass
CATEGORY = "PanoCard/convert"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"equ": ("IMAGE",),
"height": ("INT", {
"default": 832,
"min": 0,
"step": 8
}),
"order": (["nearest", "bilinear"],),
},
}
RETURN_TYPES = ("IMAGE","IMAGE","MASK")
RETURN_NAMES = ("face","long","long_mask")
OUTPUT_IS_LIST = (True, False, False)
FUNCTION = "node"
def node(self, equ, height, order):
image = equ
sub_height, sub_width = image.shape[1], image.shape[2]
if sub_height*2 != sub_width:
raise ValueError("传入图像宽高必须为2:1")
if height == 0:
height = calculate_face_heght(sub_height)
Rot = [{'yaw': 0.0, 'pitch': 0.0, 'roll': 0.0}]
equ = image.permute(0, 3, 1, 2)
long_image = py360.equi2cube(
equi = equ,
rots = Rot,
w_face = height,
mode = order,
cube_format = 'horizon')
long_image = long_image.permute(0, 2, 3, 1)
# 分割图像
face_images = []
for i in range(6):
start_idx = i * height
end_idx = (i + 1) * height
face_image = long_image[0][:, start_idx:end_idx, :]
face_images.append(face_image)
# 将分割后的图像堆叠成一个新的张量
face_images = torch.stack(face_images, dim=0) # 形状变为 [6, H, H, C]
face = face_images.unsqueeze(0)
# 添加alpha通道
if face.shape[-1] == 3:
face = add_alpha_channel(face)
if long_image.shape[-1] == 3:
long_image = add_alpha_channel(long_image)
# 简单生成MASK
mask = long_image[:, :, :, 3]
return (face, long_image, mask)
def calculate_vfov(hfov, width, height):
# 将水平视场角从度转换为弧度
hfov_rad = hfov * np.pi / 180
# 计算水平方向的半视场角
half_hfov_rad = hfov_rad / 2
# 利用宽高比计算垂直方向的半视场角的正切值
tan_half_hfov = np.tan(half_hfov_rad)
tan_half_vfov = tan_half_hfov * (height / width)
# 计算垂直方向的半视场角
half_vfov_rad = np.arctan(tan_half_vfov)
# 计算垂直视场角并转换回度
vfov = 2 * half_vfov_rad * 180 / np.pi
return vfov
class PanoImageEqu2Equ:
def __init__(self):
pass
CATEGORY = "PanoCard/convert"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"equ": ("IMAGE",),
"u_deg": ("FLOAT", {"default": 0.0,"min": -180,"max": 180,"step": 0.1}),
"v_deg": ("FLOAT", {"default": 0.0,"min": -90,"max": 90,"step": 0.1}),
"width": ("INT", {"default": 1664,"min": 0,"step": 8}),
"height": ("INT", {"default": 832,"min": 0,"step": 8}),
}
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("equ", )
FUNCTION = "node"
def node(self, equ, u_deg, v_deg, width, height):
image = equ.permute(0, 3, 1, 2)
if width == 0:
width = image.shape[3]
if height == 0:
height = image.shape[2]
# rotations
Rot = [{
'roll': 0.,
'pitch': np.pi * v_deg / 180, # rotate vertical
'yaw': np.pi * u_deg / 180, # rotate horizontal
}]
pic = py360.equi2equi(
src = image,
rots = Rot,
height = height,
width = width)
result = pic.permute(0, 2, 3, 1)
return (result,)
class PanoImagePic2Equ:
def __init__(self):
pass
CATEGORY = "PanoCard/convert"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", ),
"hfov": ("FLOAT", {"default": 60.0, "min": 0.0, "max": 180.0, "step": 0.1}),
"u_deg": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1}),
"v_deg": ("FLOAT", {"default": 0.0, "min": -90.0, "max": 90.0, "step": 0.1}),
"roll": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1}),
"height": ("INT", {"default": 1024,"min": 8,"step": 8})
},
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("equ", )
FUNCTION = "node"
def node(self, image, hfov, u_deg, v_deg, roll, height):
B = image.shape[0]
processed_images = []
x_yaw = -u_deg
y_pitch = -v_deg
for i in range(B):
# Process the image using the method
image_i = image[i]
vfov = calculate_vfov(hfov, image_i.shape[1], image_i.shape[0])
processed_image = plane_to_equirect(
image_i,
hfov,
vfov,
x_yaw,
y_pitch,
roll,
height + height
)
# Append the processed image to the list
processed_images.append(processed_image)
# Aggregate the processed images into a tensor of shape [B, H, W, C]
output_tensor = torch.stack(processed_images, dim=0)
return (output_tensor,)
class PanoImageEqu2Pic:
def __init__(self):
pass
CATEGORY = "PanoCard/convert"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"equ": ("IMAGE",),
"u_deg": ("FLOAT", {"default": 0.0,"min": -180,"max": 180,"step": 0.1}),
"v_deg": ("FLOAT", {"default": 0.0,"min": -90,"max": 90,"step": 0.1}),
"hfov": ("FLOAT", {"default": 90.0,"min": 0,"max": 180,"step": 0.1}),
"width": ("INT", {"default": 512,"min": 8,"step": 8}),
"height": ("INT", {"default": 512,"min": 8,"step": 8}),
}
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("images", )
FUNCTION = "node"
def node(self, equ, u_deg, v_deg, hfov, width, height):
image = equ.permute(0, 3, 1, 2)
# rotations
Rot = [{
'roll': 0.,
'pitch': np.pi * v_deg / 180, # rotate vertical
'yaw': np.pi * u_deg / 180, # rotate horizontal
}]
pic = py360.equi2pers(
equi = image,
rots = Rot,
fov_x = hfov,
height = height,
width = width)
result = pic.permute(0, 2, 3, 1)
return (result,)
distort.py
+500
View File
@@ -0,0 +1,500 @@
import math
import numpy as np
import torch
import torch.nn.functional as F
from .misc import *
def depolar_transform(
input_tensor: torch.Tensor,
output_size: tuple,
angle: float = 0.0,
exponent: float = 1.0,
background_color: tuple = (0, 0, 0, 0),
device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
) -> torch.Tensor:
"""
将输入图像进行反极坐标（Depolar）变换，支持透明背景和高效的插值方式。
参数:
input_tensor (torch.Tensor): 输入图像张量，形状为 [B, H, W, C]，类型为 float32。
output_size (tuple): 输出图像的尺寸 (宽度, 高度)。
angle (float): 控制扭曲的角度（以度为单位）。
exponent (float): 调整扭曲的非线性程度。
background_color (tuple): 背景颜色的RGBA值，四个整数，范围 [0, 255]。
device (torch.device): 计算设备，默认为可用的 GPU 或 CPU。
返回:
torch.Tensor: 变换后的图像，形状为 [B, H_out, W_out, C]，类型为 float32。
"""
assert input_tensor.dim() == 4, "输入张量必须是 4 维的 [B, H, W, C]"
B, H_in, W_in, C = input_tensor.shape
W_out, H_out = output_size
# 转换为 [B, C, H, W] 以适应 grid_sample 的输入格式
input_tensor = input_tensor.permute(0, 3, 1, 2).to(device) # [B, C, H_in, W_in]
# 创建输出图像的网格
# 生成两组线性空间并创建 meshgrid
theta = math.radians(90 + angle) # 转换角度到弧度
lin_w = torch.linspace(0, 1, steps=W_out, device=device)
lin_h = torch.linspace(0, 1, steps=H_out, device=device)
grid_w, grid_h = torch.meshgrid(lin_w, lin_h, indexing='ij') # [W_out, H_out]
grid_w = grid_w.t() # 转置以匹配 H_out, W_out
grid_h = grid_h.t()
# 应用指数调节非线性程度
r = grid_h ** exponent # [H_out, W_out]
# 计算极坐标(theta和r)
theta_grid = grid_w * 2 * math.pi + theta # 角度范围 [theta, theta + 2π]
# 转换为笛卡尔坐标并归一化至 [-1, 1]
x = (r * torch.cos(theta_grid)).unsqueeze(0).unsqueeze(-1) # [1, H_out, W_out, 1]
y = (r * torch.sin(theta_grid)).unsqueeze(0).unsqueeze(-1) # [1, H_out, W_out, 1]
# 将坐标从 [-1,1] 映射到 grid_sample 所需的范围
grid = torch.cat((x, y), dim=-1) # [1, H_out, W_out, 2]
grid = grid.repeat(B, 1, 1, 1) # [B, H_out, W_out, 2]
# 使用 grid_sample 进行插值
# 设定 padding_mode 为 'zeros' 以便后续可以处理背景颜色
sampled = F.grid_sample(
input_tensor,
grid,
mode='bilinear',
padding_mode='zeros',
align_corners=True
) # [B, C, H_out, W_out]
# 处理透明背景
# 创建掩码：超出输入图像区域的部分为 0，否则为 1
# 因为 grid_sample 使用 padding_mode='zeros'，所以这些区域在 sampled 中为 0
if C == 4:
alpha = sampled[:, 3:, :, :] # 获取 alpha 通道
mask = (alpha > 0).float()
else:
# 如果没有 alpha 通道，假设完全不透明
mask = torch.ones_like(sampled[:, :1, :, :])
# 转换背景颜色为 float32 并归一化
bg_color = torch.tensor(background_color, dtype=torch.float32, device=device).view(1, C, 1, 1) / 255.0
bg_color = bg_color.repeat(B, 1, H_out, W_out)
# 将 sampled 和背景颜色按掩码进行混合
output = sampled * mask + bg_color * (1 - mask)
# 转换回 [B, H_out, W_out, C] 格式
output = output.permute(0, 2, 3, 1) # [B, H_out, W_out, C]
return output.float()
'''
def plane_to_cylinder(
input_image: torch.Tensor,
fov: float = 70.0,
interpolation: str = 'nearest',
output_size: tuple = None
) -> torch.Tensor:
"""
将平面图像转换为等距圆柱投影的图像，支持指定输出尺寸、背景颜色（包括透明度），
以及选择插值方式。
参数:
input_image (torch.Tensor): 输入图像张量，形状为 [B, H, W, C]，类型为 float32，范围为 [0, 1]
fov (float): 视场角，单位为度，默认为70度
interpolation (str): 插值方式，支持 'bilinear' 和 'nearest'，默认 'bilinear'
output_size (tuple, optional): 输出图像的尺寸，格式为 (输出高度, 输出宽度)。如果未指定，则使用输入图像的尺寸
返回:
torch.Tensor: 转换后的图像张量，形状为 [B, Output_H, Output_W, C]，类型为 float32，范围为 [0, 1]
"""
# 验证插值方式
if interpolation not in ['bilinear', 'nearest']:
raise ValueError("插值方式仅支持 'bilinear' 和 'nearest'。")
# 设置输出尺寸
if output_size is not None:
if (not isinstance(output_size, tuple)) or (len(output_size) != 2):
raise ValueError("output_size 应为一个包含两个整数的元组，例如 (800, 1200)。")
output_height, output_width = output_size
else:
B, H, W, C = input_image.shape
output_height, output_width = H, W
device = input_image.device
dtype = input_image.dtype
B_in, H_in, W_in, C = input_image.shape
if C == 3:
# 添加 alpha 通道
input_image = add_alpha_channel(input_image)
# 转换输入图像形状为 [B, C, H, W]
input_image = input_image.permute(0, 3, 1, 2).to(dtype) # [B, C, H_in, W_in]
# 计算圆柱半径 R
fov_rad = math.radians(fov)
R = (W_in / 2) / math.tan(fov_rad / 2)
# 创建输出网格，确保所有操作在输入图像的设备上执行
# 创建标准化坐标网格，范围 [-1, 1]
y, x = torch.meshgrid(
torch.linspace(-1, 1, output_height, device=device, dtype=dtype),
torch.linspace(-1, 1, output_width, device=device, dtype=dtype),
indexing='ij' # 确保维度顺序正确
)
# 扩展至批次维度
x = x.unsqueeze(0).expand(B_in, -1, -1) # [B, Output_H, Output_W]
y = y.unsqueeze(0).expand(B_in, -1, -1) # [B, Output_H, Output_W]
# 计算 theta (水平角度)
# 将x从 [-1, 1] 映射到 [-W_in/2, W_in/2]
X = x * (W_in / 2)
theta = X / R # [B, Output_H, Output_W]
# 防止 theta 过大，导致 tan(theta) 过大
max_theta = math.tan(fov_rad / 2)
theta_clipped = theta.clamp(-max_theta, max_theta) # [B, Output_H, Output_W]
# 计算输入图像的 x 和 y 坐标
x_in = R * torch.tan(theta_clipped) + (W_in / 2) # [B, Output_H, Output_W]
cos_theta = torch.cos(theta_clipped) # [B, Output_H, Output_W]
cos_theta = cos_theta.clamp(min=1e-6)
y_in = y * (H_in / 2) / cos_theta + (H_in / 2) # [B, Output_H, Output_W]
# 将 [0, W_in-1] 和 [0, H_in-1] 映射到 [-1, 1]
x_norm = 2.0 * x_in / (W_in - 1) - 1.0 # [B, Output_H, Output_W]
y_norm = 2.0 * y_in / (H_in - 1) - 1.0 # [B, Output_H, Output_W]
# 创建 grid，形状为 [B, Output_H, Output_W, 2]
grid = torch.stack((x_norm, y_norm), dim=-1) # [B, Output_H, Output_W, 2]
# 使用 grid_sample 进行采样
mode = 'bilinear' if interpolation == 'bilinear' else 'nearest'
sampled = F.grid_sample(
input_image, grid, mode=mode, padding_mode='zeros', align_corners=True
) # [B, C, Output_H, Output_W]
# 将采样结果转换为 [B, Output_H, Output_W, C]
output = sampled.permute(0, 2, 3, 1) # [B, Output_H, Output_W, C]
return output # [B, Output_H, Output_W, C]
'''
def plane_to_cylinder(
input_image: torch.Tensor,
fov: float = 70.0,
interpolation: str = 'nearest',
output_size: tuple = None,
crop: bool = True
) -> torch.Tensor:
# 验证插值方式
if interpolation not in ['bilinear', 'nearest']:
raise ValueError("插值方式仅支持 'bilinear' 和 'nearest'。")
# 设置输出尺寸
if output_size is not None:
if (not isinstance(output_size, tuple)) or (len(output_size) != 2):
raise ValueError("output_size 应为一个包含两个整数的元组，例如 (800, 1200)。")
output_height, output_width = output_size
else:
B, H, W, C = input_image.shape
output_height, output_width = MUT8(H*1.2), MUT8(W*1.2)
device = input_image.device
dtype = input_image.dtype
B_in, H_in, W_in, C = input_image.shape
if C == 3:
# 添加 alpha 通道
input_image = add_alpha_channel(input_image)
# 转换输入图像形状为 [B, C, H, W]
input_image = input_image.permute(0, 3, 1, 2).to(dtype) # [B, C, H_in, W_in]
# 计算圆柱半径 R 使用绝对值fov
abs_fov = abs(fov)
abs_fov_rad = math.radians(abs_fov)
R = (W_in / 2) / math.tan(abs_fov_rad / 2)
max_theta = math.tan(abs_fov_rad / 2)
# 创建输出网格
y, x = torch.meshgrid(
torch.linspace(-1, 1, output_height, device=device, dtype=dtype),
torch.linspace(-1, 1, output_width, device=device, dtype=dtype),
indexing='ij'
)
x = x.unsqueeze(0).expand(B_in, -1, -1) # [B, Output_H, Output_W]
y = y.unsqueeze(0).expand(B_in, -1, -1) # [B, Output_H, Output_W]
# 计算 theta (水平角度)
X = x * (W_in / 2)
theta = X / R # [B, Output_H, Output_W]
# 限制 theta 范围
theta_clipped = theta.clamp(-max_theta, max_theta)
# 计算输入图像的坐标
x_in = R * torch.tan(theta_clipped) + (W_in / 2) # [B, Output_H, Output_W]
cos_theta = torch.cos(theta_clipped).clamp(min=1e-6)
# 根据fov符号调整垂直方向计算
if fov < 0:
# 动态调整垂直缩放，确保顶点不裁剪
cos_theta_min = math.cos(max_theta) # 最小cos(theta)对应最大x偏移
scale_factor = 1.0 / cos_theta_min # 基于最小cos(theta)缩放
y_in = y * (H_in / 2) * (cos_theta * scale_factor) + (H_in / 2)
else:
y_in = y * (H_in / 2) / cos_theta + (H_in / 2)
# 归一化到[-1, 1]
x_norm = 2.0 * x_in / (W_in - 1) - 1.0
y_norm = 2.0 * y_in / (H_in - 1) - 1.0
grid = torch.stack((x_norm, y_norm), dim=-1) # [B, H, W, 2]
# 采样
mode = 'bilinear' if interpolation == 'bilinear' else 'nearest'
sampled = F.grid_sample(
input_image, grid, mode=mode, padding_mode='zeros', align_corners=True
)
# 调整输出形状
output = sampled.permute(0, 2, 3, 1)
# 自动裁剪填充的零区域
if crop:
# 计算有效坐标范围
x_in_valid = (x_in >= 0) & (x_in <= W_in - 1) # [B, H, W]
y_in_valid = (y_in >= 0) & (y_in <= H_in - 1) # [B, H, W]
# 找出有效列（水平方向）
cols_valid = x_in_valid.any(dim=1) # [B, W]
left = torch.argmax(cols_valid.int(), dim=1) # [B]
reversed_cols = torch.flip(cols_valid.int(), dims=[1])
right = output_width - torch.argmax(reversed_cols, dim=1) - 1 # [B]
# 找出有效行（垂直方向）
rows_valid = y_in_valid.any(dim=2) # [B, H]
top = torch.argmax(rows_valid.int(), dim=1) # [B]
reversed_rows = torch.flip(rows_valid.int(), dims=[1])
bottom = output_height - torch.argmax(reversed_rows, dim=1) - 1 # [B]
# 假设所有batch的裁剪范围相同
left = left[0].item()
right = right[0].item()
top = top[0].item()
bottom = bottom[0].item()
print(f"裁剪范围：左={left}, 右={right}, 上={top}, 下={bottom}")
# 执行双方向裁剪
output = output[:, top:bottom+1, left:right+1, :]
return output
# 修改来自 https://github.com/willchil/ComfyUI-Environment-Visualizer
def plane_to_equirect(input_tensor, HFOV, VFOV, yaw, pitch, roll, output_width=2048):
"""
Maps an input image tensor to an equirectangular panoramic image using PyTorch tensors.
Parameters:
input_tensor (torch.Tensor): Input image as a tensor with shape [H, W, C] in RGB or RGBA format.
HFOV (float): Horizontal Field of View in degrees.
yaw (float): Yaw rotation in degrees.
pitch (float): Pitch rotation in degrees.
roll (float): Roll rotation in degrees.
output_width (int): Width of the output equirectangular image (Height will be output_width // 2).
Returns:
equirect_image (torch.Tensor): Equirectangular image tensor with shape [output_width // 2, output_width, 4].
"""
if input_tensor.ndim != 3:
raise ValueError("Input tensor must be a 3-dimensional array [H, W, C].")
H_in, W_in, C = input_tensor.shape
# Handle images with or without an alpha channel
if C == 4:
# Input has alpha channel; preserve RGB and ignore input alpha
input_image = input_tensor[:, :, :3].clone()
elif C == 3:
input_image = input_tensor.clone()
else:
raise ValueError("Input tensor must have 3 (RGB) or 4 (RGBA) channels.")
# Calculate Vertical Field of View (VFOV) to match FOV per pixel
# VFOV = HFOV * (H_in / W_in)
# Convert FOV from degrees to radians
HFOV_rad = torch.deg2rad(torch.tensor(HFOV, dtype=input_tensor.dtype, device=input_tensor.device))
VFOV_rad = torch.deg2rad(torch.tensor(VFOV, dtype=input_tensor.dtype, device=input_tensor.device))
# Compute focal lengths
fx = (W_in / 2) / torch.tan(HFOV_rad / 2)
fy = (H_in / 2) / torch.tan(VFOV_rad / 2)
# Principal point (assuming centered)
cx = W_in / 2
cy = H_in / 2
# Compute the rotation matrix from yaw, pitch, roll
def rotation_matrix(yaw, pitch, roll):
# Convert angles from degrees to radians
yaw_rad = math.radians(yaw)
pitch_rad = math.radians(pitch)
roll_rad = math.radians(roll)
# Rotation matrices around x, y, z axes
Rx = torch.tensor([
[1, 0, 0],
[0, math.cos(pitch_rad), -math.sin(pitch_rad)],
[0, math.sin(pitch_rad), math.cos(pitch_rad)]
], dtype=input_tensor.dtype, device=input_tensor.device)
Ry = torch.tensor([
[math.cos(yaw_rad), 0, math.sin(yaw_rad)],
[0, 1, 0],
[-math.sin(yaw_rad), 0, math.cos(yaw_rad)]
], dtype=input_tensor.dtype, device=input_tensor.device)
Rz = torch.tensor([
[math.cos(roll_rad), -math.sin(roll_rad), 0],
[math.sin(roll_rad), math.cos(roll_rad), 0],
[0, 0, 1]
], dtype=input_tensor.dtype, device=input_tensor.device)
# Combined rotation matrix
R = Rz @ Ry @ Rx
return R
# Compute rotation matrix and ensure it's contiguous
R = rotation_matrix(yaw, pitch, roll).T # Transpose for inverse rotation
# Equirectangular image dimensions (Height is half of Width)
W_out = output_width
H_out = output_width // 2 # Enforce 2:1 aspect ratio
# Create a meshgrid for the equirectangular image
theta = (torch.linspace(0, W_out - 1, W_out, dtype=input_tensor.dtype, device=input_tensor.device) / W_out) * 2 * math.pi - math.pi # theta from -π to π
phi = (0.5 - (torch.linspace(0, H_out - 1, H_out, dtype=input_tensor.dtype, device=input_tensor.device) / H_out)) * math.pi # phi from -π/2 to π/2
# Use torch.meshgrid with proper ordering to get [H_out, W_out]
phi_grid, theta_grid = torch.meshgrid(phi, theta, indexing='ij') # Shape: [H_out, W_out]
# Spherical to Cartesian coordinates (direction vectors)
x_s = torch.cos(phi_grid) * torch.sin(theta_grid) # Shape: [H_out, W_out]
y_s = -torch.sin(phi_grid) # Shape: [H_out, W_out]
z_s = torch.cos(phi_grid) * torch.cos(theta_grid) # Shape: [H_out, W_out]
# Stack into direction vectors
dirs = torch.stack((x_s, y_s, z_s), dim=-1) # Shape: [H_out, W_out, 3]
# Rotate direction vectors to camera coordinate system
dirs_cam = torch.matmul(dirs, R) # Shape: [H_out, W_out, 3]
dx_c, dy_c, dz_c = dirs_cam[..., 0], dirs_cam[..., 1], dirs_cam[..., 2]
# Compute valid_mask before division to avoid divide by zero
epsilon = torch.tensor(1e-6, dtype=input_tensor.dtype, device=input_tensor.device)
valid_mask = dz_c > epsilon # Points in front of the camera
# Compute x_im and y_im
x_im = (dx_c / dz_c) * fx + cx
y_im = (dy_c / dz_c) * fy + cy
# Update valid_mask with x_im and y_im in valid image range
valid_mask &= (x_im >= 0) & (x_im < W_in) & (y_im >= 0) & (y_im < H_in)
valid_mask &= torch.isfinite(x_im) & torch.isfinite(y_im)
# Prepare grid for grid_sample
# Normalize x_im and y_im to [-1, 1]
grid_x = (x_im / (W_in - 1)) * 2 - 1 # Shape: [H_out, W_out]
grid_y = (y_im / (H_in - 1)) * 2 - 1 # Shape: [H_out, W_out]
# Ensure grid_x and grid_y are within [-1, 1]
grid_x = torch.clamp(grid_x, -1.0, 1.0)
grid_y = torch.clamp(grid_y, -1.0, 1.0)
# Stack to create grid of shape [1, H_out, W_out, 2]
grid = torch.stack((grid_x, grid_y), dim=-1) # Shape: [H_out, W_out, 2]
grid = grid.unsqueeze(0) # Shape: [1, H_out, W_out, 2]
# Ensure grid is contiguous
grid = grid.contiguous()
# Prepare input image tensor
# Convert to float and permute to [C, H, W]
input_image = input_image.to(dtype=input_tensor.dtype, device=input_tensor.device).permute(2, 0, 1).unsqueeze(0) # Shape: [1, C, H_in, W_in]
# Perform remapping using grid_sample
remapped = F.grid_sample(input_image, grid, mode='bilinear', padding_mode='zeros', align_corners=True) # Shape: [1, C, H_out, W_out]
# Create an alpha channel based on valid_mask
alpha = valid_mask.unsqueeze(0).unsqueeze(0).float() # Shape: [1, 1, H_out, W_out]
# Mask the remapped RGB image with the alpha channel to set invalid regions to zero
remapped = remapped * alpha # Zero out invalid regions
# Combine the remapped RGB image with the alpha channel to create RGBA image
equirect_image = torch.cat((remapped, alpha), dim=1) # Shape: [1, C, H_out, W_out]
# Squeeze the batch dimension and permute to [H_out, W_out, C+1]
equirect_image = equirect_image.squeeze(0).permute(1, 2, 0) # Shape: [H_out, W_out, C]
equirect_image = equirect_image.clamp(0, 1)
return equirect_image
def arc_distortion(image, stretch_factor=0.5):
"""
实现垂直方向圆弧形扭曲，支持透明填充
Args:
image: 输入张量 [B,H,W,C] (C=3或4)
stretch_factor: 扭曲强度（正值向下弯曲，负值向上）
Returns:
带透明通道的扭曲图像 [B,H,W,4]
"""
B, H, W, C = image.shape
device = image.device
# 添加alpha通道
if C == 3:
alpha = torch.ones((B, H, W, 1), device=device)
img_alpha = torch.cat([image, alpha], dim=3)
else:
img_alpha = image.clone()
# 转换为BCHW格式
img = img_alpha.permute(0, 3, 1, 2) # [B, C, H, W]
# 生成归一化网格坐标
y_coord = torch.linspace(-1, 1, H, device=device)
x_coord = torch.linspace(-1, 1, W, device=device)
gy, gx = torch.meshgrid(y_coord, x_coord, indexing='ij') # [H, W]
# 垂直方向位移计算（基于水平位置）
delta_y = -stretch_factor * gx**2 # 二次函数模拟垂直弯曲
new_gy = gy + delta_y # 应用垂直位移
# 构建采样网格
grid = torch.stack([gx, new_gy], dim=-1) # [H, W, 2]
grid = grid.unsqueeze(0).expand(B, -1, -1, -1) # [B, H, W, 2]
# 执行采样（使用零填充）
distorted = F.grid_sample(
img,
grid,
mode='bilinear',
padding_mode='zeros', # 关键：超出区域填充透明
align_corners=True
)
# 恢复维度并确保alpha通道正确
output = distorted.permute(0, 2, 3, 1) # [B, H, W, C]
# 限制alpha在[0,1]范围
output[..., 3:4] = torch.clamp(output[..., 3:4], 0, 1)
return output
misc.py
+396
View File
@@ -0,0 +1,396 @@
import torch
import math
from torch import Tensor
import torchvision.transforms as t
from PIL import Image as ImageF
from PIL.Image import Image as ImageB
from typing import Literal, Any
def tensor_to_image(self):
return t.ToPILImage()(self.permute(2, 0, 1))
def image_to_tensor(self):
return t.ToTensor()(self).permute(1, 2, 0)
def get_sampler_by_name(method) -> Literal[0, 1, 2, 3, 4, 5]:
if method == "lanczos":
return ImageF.LANCZOS
elif method == "bicubic":
return ImageF.BICUBIC
elif method == "hamming":
return ImageF.HAMMING
elif method == "bilinear":
return ImageF.BILINEAR
elif method == "box":
return ImageF.BOX
elif method == "nearest":
return ImageF.NEAREST
else:
raise ValueError("Sampler not found.")
def create_rgba_image(width: int, height: int, color=(0, 0, 0, 0)) -> ImageB:
return ImageF.new("RGBA", (width, height), color)
Tensor.tensor_to_image = tensor_to_image
ImageB.image_to_tensor = image_to_tensor
# 调用的ComfyUI_ALLOR 的代码
# 创建图片
def ImageContainer(width, height, red, green, blue, alpha):
return create_rgba_image(width, height, (red, green, blue, int(alpha * 255))).image_to_tensor().unsqueeze(0)
#缩放大小
def ImageTransformResizeAbsolute(images, width, height, method):
def resize_tensor(tensor):
return tensor.tensor_to_image().resize((width, height), get_sampler_by_name(method)).image_to_tensor()
return torch.stack([resize_tensor(images[i]) for i in range(len(images))])
#缩放比例
def ImageTransformResizeRelative(images, scale_width, scale_height, method):
height, width = images[0, :, :, 0].shape
width = int(width * scale_width)
height = int(height * scale_height)
return ImageTransformResizeAbsolute(images, width, height, method)
#缩放偏移位置
def ImageTransformResizeClip(images, max_width, max_height, min_width, min_height, method):
height, width = images[0, :, :, 0].shape
if min_width >= max_width or min_height >= max_height:
return (images,)
scale_min = max(min_width / width, min_height / height)
scale_max = min(max_width / width, max_height / height)
scale = max(scale_min, scale_max)
return ImageTransformResizeRelative(images, scale, scale, method)
#两张图片绝对位置合成
def ImageCompositeAbsolute(
images_a,
images_b,
images_a_x,
images_a_y,
images_b_x,
images_b_y,
container_width,
container_height,
background,
method
):
def clip(value: float):
return value if value >= 0 else 0
# noinspection PyUnresolvedReferences
def composite(image_a, image_b):
img_a_height, img_a_width, img_a_dim = image_a.shape
img_b_height, img_b_width, img_b_dim = image_b.shape
if img_a_dim == 3:
image_a = torch.stack([
image_a[:, :, 0],
image_a[:, :, 1],
image_a[:, :, 2],
torch.ones((img_a_height, img_a_width))
], dim=2)
if img_b_dim == 3:
image_b = torch.stack([
image_b[:, :, 0],
image_b[:, :, 1],
image_b[:, :, 2],
torch.ones((img_b_height, img_b_width))
], dim=2)
container_x = max(img_a_width, img_b_width) if container_width == 0 else container_width
container_y = max(img_a_height, img_b_height) if container_height == 0 else container_height
container_a = torch.zeros((container_y, container_x, 4))
container_b = torch.zeros((container_y, container_x, 4))
img_a_height_c, img_a_width_c = [
clip((images_a_y + img_a_height) - container_y),
clip((images_a_x + img_a_width) - container_x)
]
img_b_height_c, img_b_width_c = [
clip((images_b_y + img_b_height) - container_y),
clip((images_b_x + img_b_width) - container_x)
]
if img_a_height_c <= img_a_height and img_a_width_c <= img_a_width:
container_a[
images_a_y:img_a_height + images_a_y - img_a_height_c,
images_a_x:img_a_width + images_a_x - img_a_width_c
] = image_a[
:img_a_height - img_a_height_c,
:img_a_width - img_a_width_c
]
if img_b_height_c <= img_b_height and img_b_width_c <= img_b_width:
container_b[
images_b_y:img_b_height + images_b_y - img_b_height_c,
images_b_x:img_b_width + images_b_x - img_b_width_c
] = image_b[
:img_b_height - img_b_height_c,
:img_b_width - img_b_width_c
]
if background == "images_a":
return ImageF.alpha_composite(
container_a.tensor_to_image(),
container_b.tensor_to_image()
).image_to_tensor()
else:
return ImageF.alpha_composite(
container_b.tensor_to_image(),
container_a.tensor_to_image()
).image_to_tensor()
if method == "pair":
if len(images_a) != len(images_b):
raise ValueError("Size of image_a and image_b not equals for pair batch type.")
return torch.stack([composite(images_a[i], images_b[i]) for i in range(len(images_a))])
elif method == "matrix":
return torch.stack([
composite(images_a[i], images_b[j]) for i in range(len(images_a)) for j in range(len(images_b))
])
return None
#两张图片相对位置合成
def ImageCompositeRelative(
images_a,
images_b,
images_a_x,
images_a_y,
images_b_x,
images_b_y,
background,
container_size_type,
method
):
def offset_by_percent(container_size: int, image_size: int, percent: float):
return int((container_size - image_size) * percent)
img_a_height, img_a_width = images_a[0, :, :, 0].shape
img_b_height, img_b_width = images_b[0, :, :, 0].shape
if container_size_type == "max":
container_width = max(img_a_width, img_b_width)
container_height = max(img_a_height, img_b_height)
elif container_size_type == "sum":
container_width = img_a_width + img_b_width
container_height = img_a_height + img_b_height
elif container_size_type == "sum_width":
container_width = img_a_width + img_b_width
container_height = max(img_a_height, img_b_height)
elif container_size_type == "sum_height":
container_width = max(img_a_width, img_b_width)
container_height = img_a_height + img_a_height
else:
raise ValueError()
return ImageCompositeAbsolute(
images_a,
images_b,
offset_by_percent(container_width, img_a_width, images_a_x),
offset_by_percent(container_height, img_a_height, images_a_y),
offset_by_percent(container_width, img_b_width, images_b_x),
offset_by_percent(container_height, img_b_height, images_b_y),
container_width,
container_height,
background,
method
)
def ImageTransformRotate(images, angle, expand, SSAA, method):
height, width = images[0, :, :, 0].shape
def rotate_tensor(tensor):
if method == "lanczos":
resize_sampler = ImageF.LANCZOS
rotate_sampler = ImageF.BICUBIC
elif method == "bicubic":
resize_sampler = ImageF.BICUBIC
rotate_sampler = ImageF.BICUBIC
elif method == "hamming":
resize_sampler = ImageF.HAMMING
rotate_sampler = ImageF.BILINEAR
elif method == "bilinear":
resize_sampler = ImageF.BILINEAR
rotate_sampler = ImageF.BILINEAR
elif method == "box":
resize_sampler = ImageF.BOX
rotate_sampler = ImageF.NEAREST
elif method == "nearest":
resize_sampler = ImageF.NEAREST
rotate_sampler = ImageF.NEAREST
else:
raise ValueError()
if SSAA > 1:
img = tensor.tensor_to_image()
img_us_scaled = img.resize((width * SSAA, height * SSAA), resize_sampler)
img_rotated = img_us_scaled.rotate(angle, rotate_sampler, expand == "true", fillcolor=(0, 0, 0, 0))
img_down_scaled = img_rotated.resize((img_rotated.width // SSAA, img_rotated.height // SSAA), resize_sampler)
result = img_down_scaled.image_to_tensor()
else:
img = tensor.tensor_to_image()
img_rotated = img.rotate(angle, rotate_sampler, expand == "true", fillcolor=(0, 0, 0, 0))
result = img_rotated.image_to_tensor()
return result
if angle == 0.0 or angle == 360.0:
return (images,)
else:
return torch.stack([
rotate_tensor(images[i]) for i in range(len(images))
])
def ImageTransformTranspose(images, method):
def transpose_tensor(tensor):
if method == "flip_horizontally":
transpose = ImageF.FLIP_LEFT_RIGHT
elif method == "flip_vertically":
transpose = ImageF.FLIP_TOP_BOTTOM
elif method == "rotate_90":
transpose = ImageF.ROTATE_90
elif method == "rotate_180":
transpose = ImageF.ROTATE_180
elif method == "rotate_270":
transpose = ImageF.ROTATE_270
elif method == "transpose":
transpose = ImageF.TRANSPOSE
elif method == "transverse":
transpose = ImageF.TRANSVERSE
else:
raise ValueError()
return tensor.tensor_to_image().transpose(transpose).image_to_tensor()
return torch.stack([transpose_tensor(images[i]) for i in range(len(images))])
def MUT8(width):
# 向上取整到最接近的8的倍数
adjusted_width = math.ceil(width / 8) * 8
return int(adjusted_width)
def resize_mask(mask, shape):
'''
输入mask和图片的shape，输出resize后的mask
shape: [H, W]
'''
return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])),
size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
def join_image_with_alpha(image: torch.Tensor, alpha: torch.Tensor):
batch_size = min(len(image), len(alpha))
out_images = []
alpha = 1.0 - resize_mask(alpha, image.shape[1:])
for i in range(batch_size):
out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
return torch.stack(out_images)
def add_alpha_channel(image):
# Check the shape of the input image
if image.dim() < 3:
raise ValueError("Input image must have at least 3 dimensions.")
# Get the shape of the image
shape = image.shape
# Create an alpha channel with the same shape as the last dimension except for the last dimension
alpha_channel_shape = list(shape)
alpha_channel_shape[-1] = 1
# Create the alpha channel filled with ones
alpha_channel = torch.ones(alpha_channel_shape, dtype=image.dtype, device=image.device)
# Concatenate the alpha channel to the image
image_with_alpha = torch.cat((image, alpha_channel), dim=-1)
return image_with_alpha
# 图像遮罩复合
def image_mask_composite(destination, source, x, y, mask=None, multiplier=8, resize_source=False):
destination = destination.clone().movedim(-1, 1)
source = source.clone().movedim(-1, 1)
source = source.to(destination.device)
if resize_source:
source = torch.nn.functional.interpolate(source, size=(destination.shape[2], destination.shape[3]), mode="bilinear")
x = max(-source.shape[3] * multiplier, min(x, destination.shape[3] * multiplier))
y = max(-source.shape[2] * multiplier, min(y, destination.shape[2] * multiplier))
left, top = (x // multiplier, y // multiplier)
right, bottom = (left + source.shape[3], top + source.shape[2],)
if mask is None:
mask = torch.ones_like(source)
else:
mask = mask.to(destination.device, copy=True)
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")
visible_width, visible_height = (destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),)
mask = mask[:, :, :visible_height, :visible_width]
inverse_mask = torch.ones_like(mask) - mask
source_portion = mask * source[:, :, :visible_height, :visible_width]
destination_portion = inverse_mask * destination[:, :, top:bottom, left:right]
destination[:, :, top:bottom, left:right] = source_portion + destination_portion
destination = destination.movedim(1, -1)
return destination
def crop_transparent_area(image_tensor, mode='full'):
"""
裁切掉图像中的透明部分。
参数:
image_tensor (torch.Tensor): 输入图像张量，形状为 [B, H, W, C]，类型为 float32。
mode (str): 裁剪模式，支持 'full'（全裁剪）、'horizontal'（水平裁剪）、'vertical'（垂直裁剪），默认为 'full'。
返回:
torch.Tensor: 裁剪后的图像张量，形状为 [B, H', W', C]，类型为 float32。
"""
alpha_channel = image_tensor[0, :, :, -1]
non_zero_indices = torch.nonzero(alpha_channel > 0)
if non_zero_indices.numel() == 0:
return image_tensor # 如果没有非透明部分，返回原图
min_y, min_x = non_zero_indices.min(dim=0)[0]
max_y, max_x = non_zero_indices.max(dim=0)[0]
if mode == 'full':
# 全裁剪
cropped_image = image_tensor[:, min_y:max_y+1, min_x:max_x+1, :]
elif mode == 'horizontal':
# 按照水平裁剪
cropped_image = image_tensor[:, :, min_x:max_x+1, :]
elif mode == 'vertical':
# 按照垂直裁剪
cropped_image = image_tensor[:, min_y:max_y+1, :, :]
else:
raise ValueError("裁剪模式仅支持 'full'、'horizontal' 和 'vertical'。")
return cropped_image
pano.py
+1839
View File
File diff suppressed because it is too large Load Diff
requirements.txt
+2
View File
@@ -0,0 +1,2 @@
pyequilib
opencv-python-headless
web/html/threeVisualizer.html
+33
View File
@@ -0,0 +1,33 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
<link type="text/css" rel="stylesheet" href="/extensions/ComfyUI-PanoCard/style/threeStyle.css">
<link type="text/css" rel="stylesheet" href="/extensions/ComfyUI-PanoCard/style/progressStyle.css">
</head>
<body>
<div id="progress-container">
<dialog open id="progress-dialog">
<p>
<label for="progress-indicator">加载全景图...</label>
</p>
<progress max="100" id="progress-indicator"></progress>
</dialog>
</div>
<div id="container"></div>
<script type="importmap">
{
"imports": {
"three": "/extensions/ComfyUI-PanoCard/js/three.module.js",
"three/addons/": "/extensions/ComfyUI-PanoCard/js/"
}
}
</script>
<script id="visualizer" type="module" filepath="" crossorigin src="/extensions/ComfyUI-PanoCard/js/threeVisualizer.js"></script>
</body>
</html>
web/js/OrbitControls.js
+1523
View File
File diff suppressed because it is too large Load Diff
web/js/api.js
+482
View File
@@ -0,0 +1,482 @@
class ComfyApi extends EventTarget {
#registered = new Set();
constructor() {
super();
this.api_host = location.host;
this.api_base = location.pathname.split('/').slice(0, -1).join('/');
this.initialClientId = sessionStorage.getItem("clientId");
}
apiURL(route) {
return this.api_base + route;
}
fetchApi(route, options) {
if (!options) {
options = {};
}
if (!options.headers) {
options.headers = {};
}
options.headers["Comfy-User"] = this.user;
return fetch(this.apiURL(route), options);
}
addEventListener(type, callback, options) {
super.addEventListener(type, callback, options);
this.#registered.add(type);
}
/**
* Poll status for colab and other things that don't support websockets.
*/
#pollQueue() {
setInterval(async () => {
try {
const resp = await this.fetchApi("/prompt");
const status = await resp.json();
this.dispatchEvent(new CustomEvent("status", { detail: status }));
} catch (error) {
this.dispatchEvent(new CustomEvent("status", { detail: null }));
}
}, 1000);
}
/**
* Creates and connects a WebSocket for realtime updates
* @param {boolean} isReconnect If the socket is connection is a reconnect attempt
*/
#createSocket(isReconnect) {
if (this.socket) {
return;
}
let opened = false;
let existingSession = window.name;
if (existingSession) {
existingSession = "?clientId=" + existingSession;
}
this.socket = new WebSocket(
`ws${window.location.protocol === "https:" ? "s" : ""}://${this.api_host}${this.api_base}/ws${existingSession}`
);
this.socket.binaryType = "arraybuffer";
this.socket.addEventListener("open", () => {
opened = true;
if (isReconnect) {
this.dispatchEvent(new CustomEvent("reconnected"));
}
});
this.socket.addEventListener("error", () => {
if (this.socket) this.socket.close();
if (!isReconnect && !opened) {
this.#pollQueue();
}
});
this.socket.addEventListener("close", () => {
setTimeout(() => {
this.socket = null;
this.#createSocket(true);
}, 300);
if (opened) {
this.dispatchEvent(new CustomEvent("status", { detail: null }));
this.dispatchEvent(new CustomEvent("reconnecting"));
}
});
this.socket.addEventListener("message", (event) => {
try {
if (event.data instanceof ArrayBuffer) {
const view = new DataView(event.data);
const eventType = view.getUint32(0);
const buffer = event.data.slice(4);
switch (eventType) {
case 1:
const view2 = new DataView(event.data);
const imageType = view2.getUint32(0)
let imageMime
switch (imageType) {
case 1:
default:
imageMime = "image/jpeg";
break;
case 2:
imageMime = "image/png"
}
const imageBlob = new Blob([buffer.slice(4)], { type: imageMime });
this.dispatchEvent(new CustomEvent("b_preview", { detail: imageBlob }));
break;
default:
throw new Error(`Unknown binary websocket message of type ${eventType}`);
}
}
else {
const msg = JSON.parse(event.data);
switch (msg.type) {
case "status":
if (msg.data.sid) {
this.clientId = msg.data.sid;
window.name = this.clientId; // use window name so it isnt reused when duplicating tabs
sessionStorage.setItem("clientId", this.clientId); // store in session storage so duplicate tab can load correct workflow
}
this.dispatchEvent(new CustomEvent("status", { detail: msg.data.status }));
break;
case "progress":
this.dispatchEvent(new CustomEvent("progress", { detail: msg.data }));
break;
case "executing":
this.dispatchEvent(new CustomEvent("executing", { detail: msg.data.node }));
break;
case "executed":
this.dispatchEvent(new CustomEvent("executed", { detail: msg.data }));
break;
case "execution_start":
this.dispatchEvent(new CustomEvent("execution_start", { detail: msg.data }));
break;
case "execution_success":
this.dispatchEvent(new CustomEvent("execution_success", { detail: msg.data }));
break;
case "execution_error":
this.dispatchEvent(new CustomEvent("execution_error", { detail: msg.data }));
break;
case "execution_cached":
this.dispatchEvent(new CustomEvent("execution_cached", { detail: msg.data }));
break;
default:
if (this.#registered.has(msg.type)) {
this.dispatchEvent(new CustomEvent(msg.type, { detail: msg.data }));
} else {
throw new Error(`Unknown message type ${msg.type}`);
}
}
}
} catch (error) {
console.warn("Unhandled message:", event.data, error);
}
});
}
/**
* Initialises sockets and realtime updates
*/
init() {
this.#createSocket();
}
/**
* Gets a list of extension urls
* @returns An array of script urls to import
*/
async getExtensions() {
const resp = await this.fetchApi("/extensions", { cache: "no-store" });
return await resp.json();
}
/**
* Gets a list of embedding names
* @returns An array of script urls to import
*/
async getEmbeddings() {
const resp = await this.fetchApi("/embeddings", { cache: "no-store" });
return await resp.json();
}
/**
* Loads node object definitions for the graph
* @returns The node definitions
*/
async getNodeDefs() {
const resp = await this.fetchApi("/object_info", { cache: "no-store" });
return await resp.json();
}
/**
*
* @param {number} number The index at which to queue the prompt, passing -1 will insert the prompt at the front of the queue
* @param {object} prompt The prompt data to queue
*/
async queuePrompt(number, { output, workflow }) {
const body = {
client_id: this.clientId,
prompt: output,
extra_data: { extra_pnginfo: { workflow } },
};
if (number === -1) {
body.front = true;
} else if (number != 0) {
body.number = number;
}
const res = await this.fetchApi("/prompt", {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(body),
});
if (res.status !== 200) {
throw {
response: await res.json(),
};
}
return await res.json();
}
/**
* Loads a list of items (queue or history)
* @param {string} type The type of items to load, queue or history
* @returns The items of the specified type grouped by their status
*/
async getItems(type) {
if (type === "queue") {
return this.getQueue();
}
return this.getHistory();
}
/**
* Gets the current state of the queue
* @returns The currently running and queued items
*/
async getQueue() {
try {
const res = await this.fetchApi("/queue");
const data = await res.json();
return {
// Running action uses a different endpoint for cancelling
Running: data.queue_running.map((prompt) => ({
prompt,
remove: { name: "Cancel", cb: () => api.interrupt() },
})),
Pending: data.queue_pending.map((prompt) => ({ prompt })),
};
} catch (error) {
console.error(error);
return { Running: [], Pending: [] };
}
}
/**
* Gets the prompt execution history
* @returns Prompt history including node outputs
*/
async getHistory(max_items=200) {
try {
const res = await this.fetchApi(`/history?max_items=${max_items}`);
return { History: Object.values(await res.json()) };
} catch (error) {
console.error(error);
return { History: [] };
}
}
/**
* Gets system & device stats
* @returns System stats such as python version, OS, per device info
*/
async getSystemStats() {
const res = await this.fetchApi("/system_stats");
return await res.json();
}
/**
* Sends a POST request to the API
* @param {*} type The endpoint to post to
* @param {*} body Optional POST data
*/
async #postItem(type, body) {
try {
await this.fetchApi("/" + type, {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: body ? JSON.stringify(body) : undefined,
});
} catch (error) {
console.error(error);
}
}
/**
* Deletes an item from the specified list
* @param {string} type The type of item to delete, queue or history
* @param {number} id The id of the item to delete
*/
async deleteItem(type, id) {
await this.#postItem(type, { delete: [id] });
}
/**
* Clears the specified list
* @param {string} type The type of list to clear, queue or history
*/
async clearItems(type) {
await this.#postItem(type, { clear: true });
}
/**
* Interrupts the execution of the running prompt
*/
async interrupt() {
await this.#postItem("interrupt", null);
}
/**
* Gets user configuration data and where data should be stored
* @returns { Promise<{ storage: "server" | "browser", users?: Promise<string, unknown>, migrated?: boolean }> }
*/
async getUserConfig() {
return (await this.fetchApi("/users")).json();
}
/**
* Creates a new user
* @param { string } username
* @returns The fetch response
*/
createUser(username) {
return this.fetchApi("/users", {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({ username }),
});
}
/**
* Gets all setting values for the current user
* @returns { Promise<string, unknown> } A dictionary of id -> value
*/
async getSettings() {
return (await this.fetchApi("/settings")).json();
}
/**
* Gets a setting for the current user
* @param { string } id The id of the setting to fetch
* @returns { Promise<unknown> } The setting value
*/
async getSetting(id) {
return (await this.fetchApi(`/settings/${encodeURIComponent(id)}`)).json();
}
/**
* Stores a dictionary of settings for the current user
* @param { Record<string, unknown> } settings Dictionary of setting id -> value to save
* @returns { Promise<void> }
*/
async storeSettings(settings) {
return this.fetchApi(`/settings`, {
method: "POST",
body: JSON.stringify(settings)
});
}
/**
* Stores a setting for the current user
* @param { string } id The id of the setting to update
* @param { unknown } value The value of the setting
* @returns { Promise<void> }
*/
async storeSetting(id, value) {
return this.fetchApi(`/settings/${encodeURIComponent(id)}`, {
method: "POST",
body: JSON.stringify(value)
});
}
/**
* Gets a user data file for the current user
* @param { string } file The name of the userdata file to load
* @param { RequestInit } [options]
* @returns { Promise<Response> } The fetch response object
*/
async getUserData(file, options) {
return this.fetchApi(`/userdata/${encodeURIComponent(file)}`, options);
}
/**
* Stores a user data file for the current user
* @param { string } file The name of the userdata file to save
* @param { unknown } data The data to save to the file
* @param { RequestInit & { overwrite?: boolean, stringify?: boolean, throwOnError?: boolean } } [options]
* @returns { Promise<Response> }
*/
async storeUserData(file, data, options = { overwrite: true, stringify: true, throwOnError: true }) {
const resp = await this.fetchApi(`/userdata/${encodeURIComponent(file)}?overwrite=${options?.overwrite}`, {
method: "POST",
body: options?.stringify ? JSON.stringify(data) : data,
...options,
});
if (resp.status !== 200 && options?.throwOnError !== false) {
throw new Error(`Error storing user data file '${file}': ${resp.status} ${(await resp).statusText}`);
}
return resp;
}
/**
* Deletes a user data file for the current user
* @param { string } file The name of the userdata file to delete
*/
async deleteUserData(file) {
const resp = await this.fetchApi(`/userdata/${encodeURIComponent(file)}`, {
method: "DELETE",
});
if (resp.status !== 204) {
throw new Error(`Error removing user data file '${file}': ${resp.status} ${(resp).statusText}`);
}
}
/**
* Move a user data file for the current user
* @param { string } source The userdata file to move
* @param { string } dest The destination for the file
*/
async moveUserData(source, dest, options = { overwrite: false }) {
const resp = await this.fetchApi(`/userdata/${encodeURIComponent(source)}/move/${encodeURIComponent(dest)}?overwrite=${options?.overwrite}`, {
method: "POST",
});
return resp;
}
/**
* @overload
* Lists user data files for the current user
* @param { string } dir The directory in which to list files
* @param { boolean } [recurse] If the listing should be recursive
* @param { true } [split] If the paths should be split based on the os path separator
* @returns { Promise<string[][]>> } The list of split file paths in the format [fullPath, ...splitPath]
*/
/**
* @overload
* Lists user data files for the current user
* @param { string } dir The directory in which to list files
* @param { boolean } [recurse] If the listing should be recursive
* @param { false | undefined } [split] If the paths should be split based on the os path separator
* @returns { Promise<string[]>> } The list of files
*/
async listUserData(dir, recurse, split) {
const resp = await this.fetchApi(
`/userdata?${new URLSearchParams({
recurse,
dir,
split,
})}`
);
if (resp.status === 404) return [];
if (resp.status !== 200) {
throw new Error(`Error getting user data list '${dir}': ${resp.status} ${resp.statusText}`);
}
return resp.json();
}
}
export const api = new ComfyApi();
web/js/lil-gui.module.min.js
Vendored
+8
View File
File diff suppressed because one or more lines are too long
web/js/three.module.js
+53917
View File
File diff suppressed because one or more lines are too long
web/js/threeVisualizer.js
+128
View File
@@ -0,0 +1,128 @@
import * as THREE from './three.module.js';
import { OrbitControls } from './OrbitControls.js';
import * as dat from './lil-gui.module.min.js';
import { api } from './api.js';
// 获取 DOM 元素
const visualizer = document.getElementById("visualizer");
const container = document.getElementById("container");
const progressDialog = document.getElementById("progress-dialog");
const progressIndicator = document.getElementById("progress-indicator");
// 设置 GUI 控件
const gui = new dat.GUI({ width: 150 });
const initialFov = 75;
const params = {
fov: initialFov
};
// 添加 FOV 控件
gui.add(params, 'fov', 10, 120).name('缩放').step(0.1).onChange(() => {
camera.fov = params.fov;
camera.updateProjectionMatrix();
});
gui.hide();
// 初始化渲染器
const renderer = new THREE.WebGLRenderer({
antialias: true,
extensions: { derivatives: true }
});
renderer.setPixelRatio(window.devicePixelRatio);
renderer.setSize(window.innerWidth, window.innerHeight);
container.appendChild(renderer.domElement);
// 创建 PMREM 生成器
const pmremGenerator = new THREE.PMREMGenerator(renderer);
// 创建场景
const scene = new THREE.Scene();
scene.background = new THREE.Color(0x000000);
// 添加环境光
const ambientLight = new THREE.AmbientLight(0xffffff);
scene.add(ambientLight);
// 设置相机
const camera = new THREE.PerspectiveCamera(initialFov, window.innerWidth / window.innerHeight, 0.1, 1000);
camera.position.set(0, 0, 5);
// 添加点光源
const pointLight = new THREE.PointLight(0xffffff, 15);
camera.add(pointLight);
scene.add(camera);
// 初始化轨道控制器
const controls = new OrbitControls(camera, renderer.domElement);
controls.target.set(0, 0, 0);
controls.enablePan = true;
controls.enableDamping = true;
// 处理窗口大小调整
window.addEventListener('resize', onWindowResize, false);
//显示控制
gui.show();
function onWindowResize() {
camera.aspect = window.innerWidth / window.innerHeight;
camera.updateProjectionMatrix();
renderer.setSize(window.innerWidth, window.innerHeight);
}
// 状态变量
let lastpanoImage = "";
let needUpdate = false;
// 帧更新函数
function frameUpdate() {
const panoImage = visualizer.getAttribute("pano_image");
if (panoImage === lastpanoImage) {
if (needUpdate) {
controls.update();
renderer.render(scene, camera);
}
requestAnimationFrame(frameUpdate);
} else {
needUpdate = false;
scene.clear();
progressDialog.open = true;
lastpanoImage = panoImage;
main(JSON.parse(lastpanoImage));
}
}
// 主函数：加载 pano 图像并设置场景
async function main(panoImageParams) {
let panoTexture;
console.log("加载 pano 参数:", panoImageParams);
if (panoImageParams?.filename) {
const panoImageUrl = api.apiURL('/view?' + new URLSearchParams(panoImageParams)).replace(/extensions.*\//, "");
const panoLoader = new THREE.TextureLoader();
try {
panoTexture = await panoLoader.loadAsync(panoImageUrl);
console.log("加载 pano 纹理:", panoTexture);
panoTexture.mapping = THREE.EquirectangularReflectionMapping;
panoTexture.colorSpace = THREE.SRGBColorSpace;
} catch (error) {
console.error("加载 pano 图像失败:", error);
}
}
if (panoTexture) {
scene.environment = panoTexture;
scene.background = panoTexture;
}
needUpdate = true;
progressDialog.close();
frameUpdate();
}
// 初始化调用
main();
web/style/progressStyle.css
+36
View File
@@ -0,0 +1,36 @@
dialog {
width: 100%;
text-align: center;
max-width: 20em;
color: white;
background-color: #000;
border: none;
position: relative;
transform: translate(-50%, -50%);
}
#progress-container {
position: absolute;
top: 50%;
left: 50%;
}
progress {
width: 100%;
height: 1em;
border: none;
background-color: #fff;
color: #eee;
}
progress::-webkit-progress-bar {
background-color: #333;
}
progress::-webkit-progress-value {
background-color: #eee;
}
progress::-moz-progress-bar {
background-color: #eee;
}
web/style/threeStyle.css
+91
View File
@@ -0,0 +1,91 @@
body {
margin: 0;
background-color: #000;
color: #fff;
font-family: Monospace;
font-size: 13px;
line-height: 24px;
overscroll-behavior: none;
}
a {
color: #ff0;
text-decoration: none;
}
a:hover {
text-decoration: underline;
}
button {
cursor: pointer;
text-transform: uppercase;
}
#info {
position: absolute;
top: 0px;
width: 100%;
padding: 10px;
box-sizing: border-box;
text-align: center;
-moz-user-select: none;
-webkit-user-select: none;
-ms-user-select: none;
user-select: none;
pointer-events: none;
z-index: 1; /* TODO Solve this in HTML */
}
a, button, input, select {
pointer-events: auto;
}
.lil-gui {
z-index: 2 !important; /* TODO Solve this in HTML */
}
@media all and ( max-width: 640px ) {
.lil-gui.root {
right: auto;
top: auto;
max-height: 50%;
max-width: 80%;
bottom: 0;
left: 0;
}
}
#overlay {
position: absolute;
font-size: 16px;
z-index: 2;
top: 0;
left: 0;
width: 100%;
height: 100%;
display: flex;
align-items: center;
justify-content: center;
flex-direction: column;
background: rgba(0,0,0,0.7);
}
#overlay button {
background: transparent;
border: 0;
border: 1px solid rgb(255, 255, 255);
border-radius: 4px;
color: #ffffff;
padding: 12px 18px;
text-transform: uppercase;
cursor: pointer;
}
#notSupported {
width: 50%;
margin: auto;
background-color: #f00;
margin-top: 20px;
padding: 10px;
}
web/visualization.js
+171
View File
@@ -0,0 +1,171 @@
import { app } from "../../scripts/app.js";
class Visualizer {
constructor(node, container, visualSrc) {
this.node = node;
this.iframe = document.createElement('iframe');
Object.assign(this.iframe, {
scrolling: "no",
style: {
overflow: "hidden",
border: "none",
width: "100%",
height: "100%",
display: 'none' // 默认隐藏
},
src: `/extensions/ComfyUI-PanoCard/html/${visualSrc}.html`
});
container.appendChild(this.iframe);
}
updateVisual(params) {
const iframeDocument = this.iframe.contentWindow.document;
const previewScript = iframeDocument.getElementById('visualizer');
previewScript.setAttribute("pano_image", JSON.stringify(params.pano_image));
}
remove() {
this.iframe.remove();
}
show() {
this.iframe.style.display = 'block';
}
hide() {
this.iframe.style.display = 'none';
}
}
function createVisualizer(node, inputName, typeName, app) {
node.name = inputName;
const widget = {
type: typeName,
name: "preview3d",
callback: () => {},
draw(ctx, node, widgetWidth, widgetY, widgetHeight) {
const margin = 10;
const topOffset = 35;
const leftOffset = 40;
const visible = app.canvas.ds.scale > 0.5 && this.type === typeName;
const width = widgetWidth - margin * 4;
const rect = ctx.canvas.getBoundingClientRect();
const transform = new DOMMatrix()
.scale(rect.width / ctx.canvas.width, rect.height / ctx.canvas.height)
.multiply(ctx.getTransform())
.translate(margin, margin + widgetY);
Object.assign(this.visualizer.style, {
left: `${transform.a * margin + transform.e + leftOffset}px`,
top: `${transform.d + transform.f + topOffset}px`,
width: `${width * transform.a}px`,
height: `${width * transform.d - widgetHeight - (8 * margin) * transform.d}px`,
position: "absolute",
overflow: "hidden",
zIndex: app.graph._nodes.indexOf(node),
display: visible ? 'block' : 'none'
});
Object.assign(this.visualizer.firstChild.style, {
transformOrigin: "50% 50%",
width: '100%',
height: '100%',
border: 'none'
});
}
};
const container = document.createElement('div');
container.id = `Comfy3D_${inputName}`;
container.style.display = 'none'; // 初始状态为隐藏
node.visualizer = new Visualizer(node, container, typeName);
widget.visualizer = container;
widget.parent = node;
document.body.appendChild(container);
node.addCustomWidget(widget);
node.updateParameters = (params) => {
node.visualizer.updateVisual(params);
};
node.onDrawBackground = function (ctx) {
this.visualizer.iframe.hidden = this.flags.collapsed;
};
node.onResize = function () {
let [w, h] = this.size;
w = Math.max(w, 600);
h = Math.max(h, 500);
if (w > 600) {
h = w - 100;
}
this.size = [w, h];
};
node.onRemoved = () => {
Object.values(node.widgets).forEach(widget => {
if (widget.visualizer) {
widget.visualizer.remove();
}
});
};
return { widget };
}
function registerVisualizer(nodeType, nodeData, nodeClassName, typeName) {
if (nodeData.name === nodeClassName) {
console.log(`[3D Visualizer] Registering node: ${nodeData.name}`);
const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = async function () {
const result = originalOnNodeCreated
? await originalOnNodeCreated.apply(this, arguments)
: undefined;
const existingNodes = app.graph._nodes.filter(node => node.type === nodeClassName);
const nodeName = `Preview3DNode_${existingNodes.length}`;
console.log(`[Comfy3D] Create: ${nodeName}`);
await createVisualizer(this, nodeName, typeName, app);
this.setSize([600, 500]);
return result;
};
nodeType.prototype.onExecuted = function (message) {
console.log(message);
if (message.pano_image) {
const params = { pano_image: message.pano_image[0] };
this.updateParameters(params);
this.visualizer.show();
} else {
if (this.visualizer.parentNode) {
this.visualizer.hide();
}
}
};
}
}
app.registerExtension({
name: "Gue.PanoCard.Viewer",
async init(app) {
// 初始化逻辑（如果需要）
},
async beforeRegisterNodeDef(nodeType, nodeData, app) {
registerVisualizer(nodeType, nodeData, "PanoCardViewer", "threeVisualizer");
}
});