133 lines
4.5 KiB
Python
133 lines
4.5 KiB
Python
import torch
|
|
from torchvision import transforms
|
|
from comfy import model_management
|
|
from comfy.utils import lanczos as comfy_lanczos_resize
|
|
|
|
def simple_resize(image_tensor:torch.Tensor, height, width):
|
|
"""
|
|
이미지 텐서를 지정된 높이와 너비로 단순 리사이즈합니다.
|
|
"""
|
|
transform_resize = transforms.Resize((height, width))
|
|
resized_tensor = transform_resize(image_tensor)
|
|
return resized_tensor
|
|
|
|
def prepare_mask(mask, channels):
|
|
"""
|
|
마스크를 지정된 채널 수에 맞게 확장합니다.
|
|
"""
|
|
if len(mask.shape) == 3:
|
|
mask = mask.unsqueeze(1)
|
|
mask = mask.expand(-1, channels, -1, -1)
|
|
return mask
|
|
|
|
def add_detail_transfer(target:torch.Tensor, source:torch.Tensor, blur:float, blend_ratio:float, mask:torch.Tensor=None):
|
|
"""
|
|
타겟 이미지에 소스 이미지의 디테일을 전송합니다.
|
|
|
|
Args:
|
|
target (torch.Tensor): 타겟 이미지 텐서 (B, H, W, C)
|
|
source (torch.Tensor): 소스 이미지 텐서 (B, H, W, C)
|
|
blur (float): 가우시안 블러의 시그마 값
|
|
blend_ratio (float): 블렌딩 비율
|
|
mask (torch.Tensor, optional): 마스크 텐서
|
|
|
|
Returns:
|
|
torch.Tensor: 디테일이 전송된 결과 이미지 텐서
|
|
"""
|
|
# 텐서 shape 및 디바이스 설정
|
|
B, H, W, C = target.shape
|
|
device = model_management.get_torch_device()
|
|
target_tensor = target.permute(0, 3, 1, 2).clone().to(device)
|
|
source_tensor = source.permute(0, 3, 1, 2).clone().to(device)
|
|
|
|
# 소스 이미지 리사이즈 (필요한 경우)
|
|
if target.shape[2:] != source.shape[2:]:
|
|
source_tensor = simple_resize(source_tensor, H, W)
|
|
|
|
# 배치 크기 맞추기 (필요한 경우)
|
|
if source.shape[0] < B:
|
|
source = source[0].unsqueeze(0).repeat(B, 1, 1, 1)
|
|
|
|
# 가우시안 블러 설정 및 적용
|
|
kernel_size = int(6 * int(blur) + 1)
|
|
gaussian_blur = transforms.GaussianBlur(kernel_size=(kernel_size, kernel_size), sigma=(blur, blur))
|
|
blurred_target = gaussian_blur(target_tensor)
|
|
blurred_source = gaussian_blur(source_tensor)
|
|
|
|
# 디테일 전송 수행
|
|
tensor_out = (source_tensor - blurred_source) + blurred_target
|
|
tensor_out = torch.lerp(target_tensor, tensor_out, blend_ratio)
|
|
|
|
# 마스크 적용 (있는 경우)
|
|
if mask is not None:
|
|
mask = prepare_mask(mask, C)
|
|
mask = mask.to(device)
|
|
if target.shape[2:] != mask.shape[2:]:
|
|
mask = simple_resize(mask, H, W)
|
|
tensor_out = torch.lerp(target_tensor, tensor_out, mask)
|
|
|
|
# 결과 클램핑 및 형식 변환
|
|
tensor_out = torch.clamp(tensor_out, 0, 1)
|
|
tensor_out = tensor_out.permute(0, 2, 3, 1).cpu().float()
|
|
|
|
return tensor_out
|
|
|
|
def dynamic_resize(image_tensor:torch.Tensor, max_pixels:int=1024*1024, min_pixels:int=512*512):
|
|
"""
|
|
이미지를 픽셀 수를 기준으로 동적 resize를 수행합니다.
|
|
|
|
Args:
|
|
image_tensor (torch.Tensor): 이미지 텐서
|
|
max_pixels (int): 최대 픽셀 수
|
|
min_pixels (int): 최소 픽셀 수
|
|
|
|
Returns:
|
|
torch.Tensor: resize된 이미지 텐서
|
|
"""
|
|
|
|
# 마지막 차원이 channel인지 확인
|
|
if image_tensor.shape[-1] in [3, 4]:
|
|
need_permute = True
|
|
else:
|
|
need_permute = False
|
|
|
|
if len(image_tensor.shape) == 3:
|
|
need_unsqueeze = True
|
|
else:
|
|
need_unsqueeze = False
|
|
|
|
# batch 이미지 형태로 변환
|
|
if (need_unsqueeze == False) and (need_permute == True):
|
|
img_ts = image_tensor.permute(0,3,1,2)
|
|
elif (need_unsqueeze == True) and (need_permute == True):
|
|
img_ts = image_tensor.permute(2,0,1).unsqueeze(0)
|
|
elif (need_unsqueeze == False) and (need_permute == False):
|
|
img_ts = image_tensor.unsqueeze(0)
|
|
else:
|
|
pass
|
|
|
|
_, channel, height, width = img_ts.shape
|
|
ratio = height / width
|
|
|
|
if channel != 3:
|
|
img_ts = img_ts[:,:3,:,:]
|
|
|
|
if height * width > max_pixels:
|
|
new_height = int((max_pixels * ratio)**0.5)
|
|
new_width = int(max_pixels / new_height)
|
|
elif height * width < min_pixels:
|
|
new_height = int((min_pixels * ratio)**0.5)
|
|
new_width = int(min_pixels / new_height)
|
|
else:
|
|
return image_tensor
|
|
|
|
resized_ts = comfy_lanczos_resize(img_ts, new_width, new_height)
|
|
|
|
# 원래 데이터 형태로 복구
|
|
if need_permute == True:
|
|
resized_ts = resized_ts.permute(0,2,3,1)
|
|
if need_unsqueeze == True:
|
|
resized_ts = resized_ts[0]
|
|
|
|
return resized_ts
|