194 lines
9.0 KiB
Python
194 lines
9.0 KiB
Python
"""PromptForge wrappers for the core nodes the image workflows need."""
|
|
import nodes as comfy_nodes
|
|
import node_helpers
|
|
from comfy_extras.nodes_slg import SkipLayerGuidanceDiTSimple
|
|
from comfy_extras.nodes_post_processing import ImageScaleToTotalPixels
|
|
from comfy_extras.nodes_moss_image import TextEncodeMOSSImageEdit
|
|
|
|
|
|
class PROMPTFORGE_TextEncode:
|
|
"""Positive and negative text in one node, so a workflow needs one encode instead of two."""
|
|
CATEGORY = "🔮 PromptForge/Prompts"
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("POSITIVE", "NEGATIVE")
|
|
FUNCTION = "encode"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"text": ("STRING", {"default": "", "multiline": True, "tooltip": "The prompt."}),
|
|
"negative": ("STRING", {"default": "worst quality, bad anatomy, watermark, jpeg artifact", "multiline": True, "tooltip": "What the image should avoid."}),
|
|
"use_negative": ("BOOLEAN", {"default": False, "tooltip": "Use the negative prompt. Off at low cfg (1) keeps the fast path; on raises cfg 1 to 1.5 in the sampler. At higher cfg, off makes the negative identical to the positive."}),
|
|
}
|
|
}
|
|
|
|
def encode(self, clip, text, negative, use_negative=False):
|
|
encoder = comfy_nodes.CLIPTextEncode()
|
|
positive = encoder.encode(clip, text)[0]
|
|
if not use_negative:
|
|
return positive, positive
|
|
negative_out = encoder.encode(clip, negative)[0]
|
|
return positive, node_helpers.conditioning_set_values(negative_out, {"promptforge_use_negative": True})
|
|
|
|
|
|
class PROMPTFORGE_MossEditEncode:
|
|
"""MOSS edit conditioning for both CFG passes: the prompt and the negative share the reference images."""
|
|
CATEGORY = "🔮 PromptForge/Prompts"
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("POSITIVE", "NEGATIVE")
|
|
FUNCTION = "encode"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"default": "", "multiline": True, "tooltip": "The edit instruction: what should change."}),
|
|
"negative": ("STRING", {"default": "nsfw, worst quality, bad anatomy, watermark, jpeg artifact", "multiline": True, "tooltip": "What the image should avoid. The reference images are attached to it too."}),
|
|
"edit_format": (["ai-toolkit", "ai-toolkit (kv_cache off)", "Moss-Image-Edit"], {"default": "Moss-Image-Edit"}),
|
|
"image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"image2": ("IMAGE",),
|
|
"image3": ("IMAGE",),
|
|
"image4": ("IMAGE",),
|
|
"image_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.05, "tooltip": "How hard the prompt listens to the images."}),
|
|
"negative_image_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.05}),
|
|
"use_negative": ("BOOLEAN", {"default": False, "tooltip": "Use the negative prompt. Off keeps the fast path and makes the negative identical to the positive; on raises cfg 1 to 1.5 in the sampler."}),
|
|
"vae": ("VAE",),
|
|
},
|
|
}
|
|
|
|
def encode(self, clip, prompt, negative, edit_format, image, image2=None, image3=None, image4=None,
|
|
image_strength=1.0, negative_image_strength=1.0, use_negative=False, vae=None):
|
|
images = {f"image_{index}": value for index, value in enumerate((image, image2, image3, image4)) if value is not None}
|
|
positive = TextEncodeMOSSImageEdit.execute(clip, prompt, edit_format, images, vae=vae, image_strength=image_strength).result[0]
|
|
if not use_negative:
|
|
return positive, positive
|
|
negative_out = TextEncodeMOSSImageEdit.execute(clip, negative, edit_format, images, vae=vae, image_strength=negative_image_strength).result[0]
|
|
return positive, node_helpers.conditioning_set_values(negative_out, {"promptforge_use_negative": True})
|
|
|
|
|
|
class PROMPTFORGE_KreaEditEncode:
|
|
"""Grounded Krea edit positive and negative in one node; both see the source image."""
|
|
CATEGORY = "🔮 PromptForge/Prompts"
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
|
|
RETURN_NAMES = ("POSITIVE", "NEGATIVE")
|
|
FUNCTION = "encode"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"default": "", "multiline": True, "tooltip": "The edit instruction: what should change."}),
|
|
"negative": ("STRING", {"default": "deformed, blurry, bad anatomy, bad hands, low quality", "multiline": True}),
|
|
"image": ("IMAGE",),
|
|
"use_negative": ("BOOLEAN", {"default": False, "tooltip": "Use the negative prompt. Off keeps the fast path and makes the negative identical to the positive; on raises cfg 1 to 1.5 in the sampler."}),
|
|
},
|
|
"optional": {
|
|
"grounding_px": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 64}),
|
|
"system_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
def encode(self, clip, prompt, negative, image, use_negative=False, grounding_px=0, system_prompt=""):
|
|
grounded = comfy_nodes.NODE_CLASS_MAPPINGS["Krea2EditGroundedEncode"]()
|
|
positive = grounded.encode(clip, prompt, image=image, grounding_px=grounding_px, system_prompt=system_prompt)[0]
|
|
if not use_negative:
|
|
return positive, positive
|
|
negative_out = comfy_nodes.NODE_CLASS_MAPPINGS["Krea2EditGroundedEncode"]().encode(
|
|
clip, negative, image=image, grounding_px=grounding_px, system_prompt=system_prompt)[0]
|
|
return positive, node_helpers.conditioning_set_values(negative_out, {"promptforge_use_negative": True})
|
|
|
|
|
|
class PROMPTFORGE_SkipLayerGuidance:
|
|
CATEGORY = "🔮 PromptForge/Sampling"
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"single_layers": ("STRING", {"default": "22"}),
|
|
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
}
|
|
}
|
|
|
|
def patch(self, model, single_layers, start_percent, end_percent):
|
|
return SkipLayerGuidanceDiTSimple.execute(model, start_percent, end_percent, "", single_layers).result
|
|
|
|
|
|
class PROMPTFORGE_ImageScale:
|
|
CATEGORY = "🔮 PromptForge/Load"
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT")
|
|
RETURN_NAMES = ("IMAGE", "WIDTH", "HEIGHT")
|
|
FUNCTION = "scale"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"megapixels": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 16.0, "step": 0.01}),
|
|
"method": (["lanczos", "area", "bilinear", "bicubic", "nearest-exact"], {"default": "lanczos"}),
|
|
"step": ("INT", {"default": 16, "min": 1, "max": 256}),
|
|
}
|
|
}
|
|
|
|
def scale(self, image, megapixels, method, step):
|
|
out = ImageScaleToTotalPixels.execute(image, method, megapixels, step).result[0]
|
|
return (out, int(out.shape[2]), int(out.shape[1]))
|
|
|
|
|
|
class PROMPTFORGE_VAEDecode:
|
|
CATEGORY = "🔮 PromptForge/Load"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "decode"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"samples": ("LATENT",), "vae": ("VAE",)}}
|
|
|
|
def decode(self, samples, vae):
|
|
return comfy_nodes.VAEDecode().decode(vae, samples)
|
|
|
|
|
|
class PROMPTFORGE_SaveImage:
|
|
CATEGORY = "🔮 PromptForge/Load"
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "save"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"images": ("IMAGE",), "filename_prefix": ("STRING", {"default": "PromptForge/image"})}}
|
|
|
|
def save(self, images, filename_prefix="PromptForge/image"):
|
|
return comfy_nodes.NODE_CLASS_MAPPINGS["SaveImage"]().save_images(images, filename_prefix=filename_prefix)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PROMPTFORGE_TextEncode": PROMPTFORGE_TextEncode,
|
|
"PROMPTFORGE_MossEditEncode": PROMPTFORGE_MossEditEncode,
|
|
"PROMPTFORGE_KreaEditEncode": PROMPTFORGE_KreaEditEncode,
|
|
"PROMPTFORGE_SkipLayerGuidance": PROMPTFORGE_SkipLayerGuidance,
|
|
"PROMPTFORGE_ImageScale": PROMPTFORGE_ImageScale,
|
|
"PROMPTFORGE_VAEDecode": PROMPTFORGE_VAEDecode,
|
|
"PROMPTFORGE_SaveImage": PROMPTFORGE_SaveImage,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PROMPTFORGE_TextEncode": "PromptForge Text Encode",
|
|
"PROMPTFORGE_MossEditEncode": "PromptForge MOSS Edit Encode",
|
|
"PROMPTFORGE_KreaEditEncode": "PromptForge Krea Edit Encode",
|
|
"PROMPTFORGE_SkipLayerGuidance": "PromptForge Skip Layer Guidance",
|
|
"PROMPTFORGE_ImageScale": "PromptForge Image Scale",
|
|
"PROMPTFORGE_VAEDecode": "PromptForge VAE Decode",
|
|
"PROMPTFORGE_SaveImage": "PromptForge Save Image",
|
|
}
|