66 lines
2.9 KiB
Python
66 lines
2.9 KiB
Python
import torch
|
|
import os
|
|
import folder_paths
|
|
|
|
class T2I_StyleApply:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"conditioning": ("CONDITIONING", ),
|
|
"image": ("IMAGE", ),
|
|
"style_model": (folder_paths.get_filename_list("style_models"), ),
|
|
"clip_vision": (folder_paths.get_filename_list("clip_vision"), ),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}),
|
|
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}}
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "apply_stylemodel"
|
|
DESCRIPTION = "T2I StyleApply with integrated loaders, strength and start/end percent control"
|
|
|
|
def apply_stylemodel(self, conditioning, image, style_model, clip_vision, strength=1.0, start_percent=0.0, end_percent=1.0):
|
|
from comfy.sd import load_style_model
|
|
from comfy.clip_vision import load as load_clip_vision
|
|
|
|
# Load CLIP Vision model
|
|
clip_vision_path = folder_paths.get_full_path("clip_vision", clip_vision)
|
|
clip_vision_model = load_clip_vision(clip_vision_path)
|
|
|
|
# Encode image with CLIP Vision
|
|
clip_vision_output = clip_vision_model.encode_image(image)
|
|
|
|
# Load Style Model
|
|
style_model_path = folder_paths.get_full_path("style_models", style_model)
|
|
style_model_loaded = load_style_model(style_model_path)
|
|
|
|
# Apply style
|
|
cond = style_model_loaded.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
|
|
cond = strength * cond
|
|
|
|
c = []
|
|
for t in conditioning:
|
|
# Create conditioning with style for the specified range
|
|
n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
|
|
n[1]["start_percent"] = start_percent
|
|
n[1]["end_percent"] = end_percent
|
|
c.append(n)
|
|
|
|
# Add original conditioning for outside the range
|
|
# Before start_percent
|
|
if start_percent > 0.0:
|
|
n_before = [t[0], t[1].copy()]
|
|
n_before[1]["start_percent"] = 0.0
|
|
n_before[1]["end_percent"] = start_percent
|
|
c.append(n_before)
|
|
|
|
# After end_percent
|
|
if end_percent < 1.0:
|
|
n_after = [t[0], t[1].copy()]
|
|
n_after[1]["start_percent"] = end_percent
|
|
n_after[1]["end_percent"] = 1.0
|
|
c.append(n_after)
|
|
|
|
return (c, )
|
|
|
|
NODE_CLASS_MAPPINGS = {"T2I_StyleApply": T2I_StyleApply}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"T2I_StyleApply": "T2I Style Apply Advanced"} |