111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
import torch
|
|
import os
|
|
import numpy as np
|
|
import folder_paths
|
|
import comfy.model_management
|
|
from comfy.clip_vision import load_torch_file, clip_preprocess
|
|
|
|
def get_model_path(model_name):
|
|
try:
|
|
path = folder_paths.get_folder_paths(model_name)[0]
|
|
except:
|
|
path = os.path.join(folder_paths.models_dir, model_name)
|
|
return path
|
|
|
|
class ClipVisionStyleLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
crop_methods = ["none", "center", "mask"]
|
|
|
|
return {"required":
|
|
{"clip_vision": (folder_paths.get_filename_list("clip_vision"), ),
|
|
"style_model": (folder_paths.get_filename_list("style_models"), ),
|
|
"image": ("IMAGE",),
|
|
"crop_method": (crop_methods, {"default": "none"}),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK", ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STYLE_MODEL", "CLIP_VISION_OUTPUT")
|
|
RETURN_NAMES = ("IMAGE", "STYLE_MODEL", "CLIP_VISION_OUTPUT")
|
|
OUTPUT_IS_LIST = (True, False, False)
|
|
FUNCTION = "process_image"
|
|
CATEGORY = "🧪AILab/⚛️ReduxFineTune"
|
|
|
|
def crop_center(self, image, target_size=224):
|
|
h, w = image.shape[1], image.shape[2]
|
|
crop_size = min(h, w)
|
|
y = (h - crop_size) // 2
|
|
x = (w - crop_size) // 2
|
|
cropped = image[:, y:y+crop_size, x:x+crop_size, :]
|
|
return cropped
|
|
|
|
def crop_mask(self, image, mask, target_size=None):
|
|
if mask.shape[1:3] != image.shape[1:3]:
|
|
mask = torch.nn.functional.interpolate(
|
|
mask.unsqueeze(1),
|
|
size=(image.shape[1], image.shape[2]),
|
|
mode='bilinear',
|
|
antialias=False
|
|
).squeeze(1)
|
|
|
|
if mask.device.type != 'cpu':
|
|
mask_np = mask[0].cpu().numpy()
|
|
else:
|
|
mask_np = mask[0].numpy()
|
|
|
|
y_indices, x_indices = np.where(mask_np > 0.05)
|
|
if len(y_indices) == 0 or len(x_indices) == 0:
|
|
return self.crop_center(image)
|
|
|
|
y1, y2 = np.min(y_indices), np.max(y_indices)
|
|
x1, x2 = np.min(x_indices), np.max(x_indices)
|
|
|
|
cropped = image[:, y1:y2, x1:x2, :]
|
|
return cropped
|
|
|
|
def process_image(self, clip_vision, style_model, image, crop_method, mask=None):
|
|
# Load CLIP model
|
|
clip_path = folder_paths.get_full_path("clip_vision", clip_vision)
|
|
clip_vision = comfy.clip_vision.load(clip_path)
|
|
if clip_vision is None:
|
|
raise Exception(f"CLIP Vision model {clip_vision} not found")
|
|
|
|
# Load style model
|
|
style_model_path = folder_paths.get_full_path("style_models", style_model)
|
|
style_model = comfy.sd.load_style_model(style_model_path)
|
|
|
|
# Process image based on crop method
|
|
if crop_method == "none":
|
|
cropped_image = image
|
|
crop_flag = False
|
|
elif crop_method == "center":
|
|
cropped_image = self.crop_center(image)
|
|
crop_flag = True
|
|
elif crop_method == "mask" and mask is not None:
|
|
cropped_image = self.crop_mask(image, mask)
|
|
crop_flag = True
|
|
else:
|
|
cropped_image = image
|
|
crop_flag = False
|
|
|
|
# Clone image for CLIP processing
|
|
clip_image = cropped_image.clone()
|
|
|
|
# Encode with CLIP
|
|
try:
|
|
out = clip_vision.encode_image(clip_image, crop=crop_flag)
|
|
except:
|
|
out = clip_vision.encode_image(clip_image)
|
|
|
|
return [cropped_image], style_model, out
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ClipVisionStyleLoader": ClipVisionStyleLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ClipVisionStyleLoader": "CLIP Vision + Style Loader",
|
|
} |