Spaces:
Running
Running
File size: 3,862 Bytes
028694a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
from torch import Tensor
from nodes import VAEEncode
import comfy.utils
from comfy.sd import VAE
from .control_reference import ReferenceAdvanced, ReferenceOptions, ReferenceType, ReferencePreprocWrapper
# node for ReferenceCN
class ReferenceControlNetNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"reference_type": (ReferenceType._LIST,),
"style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"ref_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
},
}
RETURN_TYPES = ("CONTROL_NET", )
FUNCTION = "load_controlnet"
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/Reference"
def load_controlnet(self, reference_type: str, style_fidelity: float, ref_weight: float):
ref_opts = ReferenceOptions.create_combo(reference_type=reference_type, style_fidelity=style_fidelity, ref_weight=ref_weight)
controlnet = ReferenceAdvanced(ref_opts=ref_opts, timestep_keyframes=None)
return (controlnet,)
class ReferenceControlFinetune:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"attn_style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"attn_ref_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"attn_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"adain_style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"adain_ref_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
"adain_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
},
}
RETURN_TYPES = ("CONTROL_NET", )
FUNCTION = "load_controlnet"
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/Reference"
def load_controlnet(self,
attn_style_fidelity: float, attn_ref_weight: float, attn_strength: float,
adain_style_fidelity: float, adain_ref_weight: float, adain_strength: float):
ref_opts = ReferenceOptions(reference_type=ReferenceType.ATTN_ADAIN,
attn_style_fidelity=attn_style_fidelity, attn_ref_weight=attn_ref_weight, attn_strength=attn_strength,
adain_style_fidelity=adain_style_fidelity, adain_ref_weight=adain_ref_weight, adain_strength=adain_strength)
controlnet = ReferenceAdvanced(ref_opts=ref_opts, timestep_keyframes=None)
return (controlnet,)
class ReferencePreprocessorNode:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE", ),
"vae": ("VAE", ),
"latent_size": ("LATENT", ),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("proc_IMAGE",)
FUNCTION = "preprocess_images"
CATEGORY = "Adv-ControlNet 🛂🅐🅒🅝/Reference/preprocess"
def preprocess_images(self, vae: VAE, image: Tensor, latent_size: Tensor):
# first, resize image to match latents
image = image.movedim(-1,1)
image = comfy.utils.common_upscale(image, latent_size["samples"].shape[3] * 8, latent_size["samples"].shape[2] * 8, 'nearest-exact', "center")
image = image.movedim(1,-1)
# then, vae encode
try:
image = vae.vae_encode_crop_pixels(image)
except Exception:
image = VAEEncode.vae_encode_crop_pixels(image)
encoded = vae.encode(image[:,:,:,:3])
return (ReferencePreprocWrapper(condhint=encoded),)
|