129 lines
4.8 KiB
Python
129 lines
4.8 KiB
Python
import comfy.utils
|
|
import comfy.sd
|
|
import folder_paths
|
|
import node_helpers
|
|
from nodes import MAX_RESOLUTION
|
|
from PIL import Image, ImageOps, ImageSequence, ImageFile
|
|
import numpy as np
|
|
import torch
|
|
|
|
class CheckpointLoaderText:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"ckpt_name": ("STRING", {"multiline": True, "dynamicPrompts": False}),
|
|
}}
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "TextFormatter/BasicNodes"
|
|
def load_checkpoint(self, ckpt_name):
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
|
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
return out[:3]
|
|
|
|
class LoraLoaderText:
|
|
def __init__(self):
|
|
self.loaded_lora = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "model": ("MODEL",),
|
|
"clip": ("CLIP", ),
|
|
"lora_name": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
|
}}
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
FUNCTION = "load_lora"
|
|
|
|
CATEGORY = "loaders"
|
|
CATEGORY = "TextFormatter/BasicNodes"
|
|
def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
|
|
if strength_model == 0 and strength_clip == 0:
|
|
return (model, clip)
|
|
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
lora = None
|
|
if self.loaded_lora is not None:
|
|
if self.loaded_lora[0] == lora_path:
|
|
lora = self.loaded_lora[1]
|
|
else:
|
|
temp = self.loaded_lora
|
|
self.loaded_lora = None
|
|
del temp
|
|
|
|
if lora is None:
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
self.loaded_lora = (lora_path, lora)
|
|
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
|
|
return (model_lora, clip_lora)
|
|
|
|
class EmptyLatentImageText:
|
|
def __init__(self):
|
|
self.device = comfy.model_management.intermediate_device()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
|
|
"height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
|
|
"batch_size": ("STRING", {"multiline": True, "dynamicPrompts": True})
|
|
#"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."})
|
|
}
|
|
}
|
|
RETURN_TYPES = ("LATENT",)
|
|
OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "TextFormatter/BasicNodes"
|
|
DESCRIPTION = "Create a new batch of empty latent images to be denoised via sampling."
|
|
|
|
def generate(self, width, height, batch_size=1):
|
|
latent = torch.zeros([int(batch_size), 4, height // 8, width // 8], device=self.device)
|
|
return ({"samples":latent}, )
|
|
|
|
class LoadImageText:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"path": ("STRING", {"multiline": True, "dynamicPrompts": True})},
|
|
}
|
|
|
|
CATEGORY = "TextFormatter/BasicNodes"
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "load_image"
|
|
def load_image(self, path):
|
|
img = node_helpers.pillow(Image.open, path)
|
|
|
|
output_images = []
|
|
w, h = None, None
|
|
|
|
excluded_formats = ['MPO']
|
|
|
|
for i in ImageSequence.Iterator(img):
|
|
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
|
|
|
if i.mode == 'I':
|
|
i = i.point(lambda i: i * (1 / 255))
|
|
image = i.convert("RGB")
|
|
|
|
if len(output_images) == 0:
|
|
w = image.size[0]
|
|
h = image.size[1]
|
|
|
|
if image.size[0] != w or image.size[1] != h:
|
|
continue
|
|
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
output_images.append(image)
|
|
|
|
if len(output_images) > 1 and img.format not in excluded_formats:
|
|
output_image = torch.cat(output_images, dim=0)
|
|
else:
|
|
output_image = output_images[0]
|
|
|
|
return (output_image,) |