134 lines
4.9 KiB
Python
134 lines
4.9 KiB
Python
"""PromptForge wrappers for the core nodes the image workflows need."""
|
|
import nodes as comfy_nodes
|
|
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:
|
|
CATEGORY = "🔮 PromptForge/Prompts"
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"clip": ("CLIP",), "text": ("STRING", {"default": "", "multiline": True})}}
|
|
|
|
def encode(self, clip, text):
|
|
return comfy_nodes.CLIPTextEncode().encode(clip, text)
|
|
|
|
|
|
class PROMPTFORGE_MossEditEncode:
|
|
CATEGORY = "🔮 PromptForge/Prompts"
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"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}),
|
|
"vae": ("VAE",),
|
|
},
|
|
}
|
|
|
|
def encode(self, clip, prompt, edit_format, image, image2=None, image3=None, image4=None, image_strength=1.0, vae=None):
|
|
images = {f"image_{index}": value for index, value in enumerate((image, image2, image3, image4)) if value is not None}
|
|
return TextEncodeMOSSImageEdit.execute(clip, prompt, edit_format, images, vae=vae, image_strength=image_strength).result
|
|
|
|
|
|
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_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_SkipLayerGuidance": "PromptForge Skip Layer Guidance",
|
|
"PROMPTFORGE_ImageScale": "PromptForge Image Scale",
|
|
"PROMPTFORGE_VAEDecode": "PromptForge VAE Decode",
|
|
"PROMPTFORGE_SaveImage": "PromptForge Save Image",
|
|
}
|