🎉 init: jannchie's ComfyUI custom nodes
This commit is contained in:
Jianqi Pan
+505
@@ -0,0 +1,505 @@
|
||||
import contextlib
|
||||
import random
|
||||
from collections import Counter
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from compel import Compel, DiffusersTextualInversionManager
|
||||
from diffusers import StableDiffusionPipeline
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
import comfy.model_management
|
||||
import folder_paths
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
from .pipelines import PipelineWrapper, schedulers
|
||||
|
||||
|
||||
def get_prompt_embeds(pipe, prompt, negative_prompt):
|
||||
textual_inversion_manager = DiffusersTextualInversionManager(pipe)
|
||||
compel = Compel(
|
||||
tokenizer=pipe.tokenizer,
|
||||
text_encoder=pipe.text_encoder,
|
||||
textual_inversion_manager=textual_inversion_manager,
|
||||
truncate_long_prompts=False,
|
||||
)
|
||||
|
||||
prompt_embeds = compel.build_conditioning_tensor(prompt)
|
||||
negative_prompt_embeds = compel.build_conditioning_tensor(negative_prompt)
|
||||
[
|
||||
prompt_embeds,
|
||||
negative_prompt_embeds,
|
||||
] = compel.pad_conditioning_tensors_to_same_length(
|
||||
[prompt_embeds, negative_prompt_embeds]
|
||||
)
|
||||
return prompt_embeds, negative_prompt_embeds
|
||||
|
||||
|
||||
def latents_to_tensor(pipeline, latents):
|
||||
image_numpy = pipeline.decode_latents(latents) # numpy
|
||||
return torch.tensor(image_numpy)
|
||||
|
||||
|
||||
def prepare_latents(
|
||||
pipe: StableDiffusionPipeline,
|
||||
batch_size: int,
|
||||
num_channels_latents: int,
|
||||
height: int,
|
||||
width: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
generator: torch.Generator,
|
||||
latents=None,
|
||||
):
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
height // pipe.vae_scale_factor,
|
||||
width // pipe.vae_scale_factor,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * pipe.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
|
||||
class GetFilledColorImage:
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "Jannchie"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"width": (
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 0,
|
||||
"max": 8192,
|
||||
"step": 64,
|
||||
"display": "number",
|
||||
},
|
||||
),
|
||||
"height": (
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 0,
|
||||
"max": 8192,
|
||||
"step": 64,
|
||||
"display": "number",
|
||||
},
|
||||
),
|
||||
"red": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.1,
|
||||
"display": "number",
|
||||
},
|
||||
),
|
||||
"green": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.1,
|
||||
"display": "number",
|
||||
},
|
||||
),
|
||||
"blue": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.1,
|
||||
"display": "number",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def run(self, width, height, red, green, blue):
|
||||
image = torch.tensor(np.full((height, width, 3), (red, green, blue)))
|
||||
# 再转换成 0 - 1 之间的浮点数
|
||||
image = image
|
||||
image = image.unsqueeze(0)
|
||||
return (image,)
|
||||
|
||||
|
||||
class DiffusersCompelPromptEmbedding:
|
||||
CATEGORY = "Jannchie"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("DIFFUSERS_PROMPT_EMBEDDING", "DIFFUSERS_PROMPT_EMBEDDING")
|
||||
RETURN_NAMES = ("positive prompt embedding", "negative prompt embedding")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pipeline": ("DIFFUSERS_PIPELINE",),
|
||||
"positive_prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "(masterpiece)1.2, (best quality)1.4",
|
||||
},
|
||||
),
|
||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
}
|
||||
}
|
||||
|
||||
def run(
|
||||
self,
|
||||
pipeline: StableDiffusionPipeline,
|
||||
positive_prompt: str,
|
||||
negative_prompt: str,
|
||||
):
|
||||
return get_prompt_embeds(pipeline, positive_prompt, negative_prompt)
|
||||
|
||||
|
||||
class DiffusersTextureInversionLoader:
|
||||
CATEGORY = "Jannchie"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("DIFFUSERS_PIPELINE",)
|
||||
RETURN_NAMES = ("pipeline",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pipeline": ("DIFFUSERS_PIPELINE",),
|
||||
"texture_inversion": (folder_paths.get_filename_list("embeddings"),),
|
||||
},
|
||||
}
|
||||
|
||||
def run(self, pipeline: StableDiffusionPipeline, texture_inversion: str):
|
||||
with contextlib.suppress(Exception):
|
||||
path = folder_paths.get_full_path("embeddings", texture_inversion)
|
||||
token = texture_inversion.split(".")[0]
|
||||
pipeline.load_textual_inversion(path, token=token)
|
||||
return (pipeline,)
|
||||
|
||||
|
||||
class GetAverageColorFromImage:
|
||||
CATEGORY = "Jannchie"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("FLOAT", "FLOAT", "FLOAT")
|
||||
RETURN_NAMES = ("red", "green", "blue")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"average": ("STRING", {"default": "mean", "options": ["mean", "mode"]}),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK",),
|
||||
},
|
||||
}
|
||||
|
||||
def run(self, image: torch.Tensor, average: str, mask: torch.Tensor = None):
|
||||
if average == "mean":
|
||||
return self.run_avg(image, mask)
|
||||
elif average == "mode":
|
||||
return self.run_mode(image, mask)
|
||||
|
||||
def run_avg(self, image: torch.Tensor, mask: torch.Tensor = None):
|
||||
if mask is not None:
|
||||
mask = mask.unsqueeze(1)
|
||||
masked_image = image * mask if mask is not None else image
|
||||
pixel_sum = torch.sum(masked_image, dim=(2, 3))
|
||||
pixel_count = (
|
||||
torch.sum(mask, dim=(2, 3))
|
||||
if mask is not None
|
||||
else torch.prod(torch.tensor(image.shape[2:]))
|
||||
)
|
||||
average_rgb = pixel_sum / pixel_count.unsqueeze(1)
|
||||
|
||||
average_rgb = torch.round(average_rgb)
|
||||
|
||||
return tuple(average_rgb.squeeze().tolist())
|
||||
|
||||
def run_mode(self, image: torch.Tensor, mask: torch.Tensor = None):
|
||||
image = image.permute(0, 3, 1, 2)
|
||||
if mask is not None:
|
||||
mask = mask.unsqueeze(1)
|
||||
|
||||
masked_image = image * mask if mask is not None else image
|
||||
pixel_values = masked_image.view(
|
||||
masked_image.shape[0], masked_image.shape[1], -1
|
||||
)
|
||||
pixel_values = pixel_values.permute(0, 2, 1)
|
||||
pixel_values = pixel_values.reshape(-1, pixel_values.shape[2])
|
||||
pixel_values = [
|
||||
tuple(color.tolist()) for color in pixel_values.numpy() if color.max() > 0
|
||||
]
|
||||
|
||||
if not pixel_values:
|
||||
return (0, 0, 0)
|
||||
|
||||
color_counts = Counter(pixel_values)
|
||||
|
||||
return max(color_counts, key=color_counts.get)
|
||||
|
||||
|
||||
class DiffusersPipeline:
|
||||
CATEGORY = "Jannchie"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("DIFFUSERS_PIPELINE",)
|
||||
RETURN_NAMES = ("pipeline",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
||||
},
|
||||
"optional": {
|
||||
"vae_name": (
|
||||
folder_paths.get_filename_list("vae") + ["-"],
|
||||
{"default": "-"},
|
||||
),
|
||||
"scheduler_name": (
|
||||
list(schedulers.keys()) + ["-"],
|
||||
{
|
||||
"default": "-",
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def run(self, ckpt_name: str, vae_name: str = None, scheduler_name: str = None):
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
if vae_name == "-":
|
||||
vae_path = None
|
||||
else:
|
||||
vae_path = folder_paths.get_full_path("vae", vae_name)
|
||||
if scheduler_name == "-":
|
||||
scheduler_name = None
|
||||
|
||||
self.pipeline_wrapper = PipelineWrapper(ckpt_path, vae_path, scheduler_name)
|
||||
return (self.pipeline_wrapper.pipeline,)
|
||||
|
||||
|
||||
class DiffusersPrepareLatents:
|
||||
CATEGORY = "Jannchie"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("latents",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pipeline": ("DIFFUSERS_PIPELINE",),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}),
|
||||
"height": ("INT", {"default": 512, "min": 0, "max": 8192, "step": 64}),
|
||||
"width": ("INT", {"default": 512, "min": 0, "max": 8192, "step": 64}),
|
||||
},
|
||||
"optional": {
|
||||
"latents": ("LATENT", {"default": None}),
|
||||
"seed": (
|
||||
"INT",
|
||||
{"default": None, "min": 0, "step": 1, "max": 999999999},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def run(
|
||||
self,
|
||||
pipeline: StableDiffusionPipeline,
|
||||
batch_size: int = 1,
|
||||
height: int = 512,
|
||||
width: int = 512,
|
||||
latents: torch.Tensor | None = None,
|
||||
seed: int | None = None,
|
||||
):
|
||||
if seed is None:
|
||||
seed = random.randint(0, 999999999)
|
||||
device = comfy.model_management.get_torch_device()
|
||||
generator = torch.Generator(device)
|
||||
generator.manual_seed(seed)
|
||||
latents = prepare_latents(
|
||||
pipe=pipeline,
|
||||
batch_size=batch_size,
|
||||
num_channels_latents=4,
|
||||
height=height,
|
||||
width=width,
|
||||
dtype=comfy.model_management.VAE_DTYPE,
|
||||
device=device,
|
||||
generator=generator,
|
||||
latents=latents,
|
||||
)
|
||||
return (latents,)
|
||||
|
||||
|
||||
class DiffusersDecoder:
|
||||
CATEGORY = "Jannchie"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pipeline": ("DIFFUSERS_PIPELINE",),
|
||||
"latents": ("LATENT",),
|
||||
},
|
||||
}
|
||||
|
||||
def run(self, pipeline: StableDiffusionPipeline, latents: torch.Tensor):
|
||||
return (latents_to_tensor(pipeline, latents),)
|
||||
|
||||
|
||||
class DiffusersGenerate:
|
||||
CATEGORY = "Jannchie"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pipeline": ("DIFFUSERS_PIPELINE",),
|
||||
"positive_prompt_embedding": ("DIFFUSERS_PROMPT_EMBEDDING",),
|
||||
"negative_prompt_embedding": ("DIFFUSERS_PROMPT_EMBEDDING",),
|
||||
"strength": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.02},
|
||||
),
|
||||
"num_inference_steps": (
|
||||
"INT",
|
||||
{"default": 30, "min": 1, "max": 100, "step": 1},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"images": ("IMAGE",),
|
||||
"seed": (
|
||||
"INT",
|
||||
{"default": 0, "min": 0, "step": 1, "max": 999999999999},
|
||||
),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}),
|
||||
"width": (
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 8192,
|
||||
"step": 64,
|
||||
},
|
||||
),
|
||||
"height": (
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 8192,
|
||||
"step": 64,
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def run(
|
||||
self,
|
||||
pipeline: StableDiffusionPipeline,
|
||||
positive_prompt_embedding: torch.Tensor,
|
||||
negative_prompt_embedding: torch.Tensor,
|
||||
width: int,
|
||||
height: int,
|
||||
batch_size: int,
|
||||
images: torch.Tensor | None = None,
|
||||
num_inference_steps: int = 30,
|
||||
strength: float = 1.0,
|
||||
seed=None,
|
||||
):
|
||||
pbar = ProgressBar(int(num_inference_steps * strength))
|
||||
device = comfy.model_management.get_torch_device()
|
||||
if not seed:
|
||||
seed = random.randint(0, 999999999999)
|
||||
generator = torch.Generator(device)
|
||||
generator.manual_seed(seed)
|
||||
# (B, H, W, C) to (B, C, H, W)
|
||||
if images is None:
|
||||
latents = prepare_latents(
|
||||
pipe=pipeline,
|
||||
batch_size=batch_size,
|
||||
num_channels_latents=4,
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
dtype=comfy.model_management.VAE_DTYPE,
|
||||
device=device,
|
||||
)
|
||||
images = latents_to_tensor(pipeline, latents).permute(0, 3, 1, 2)
|
||||
else:
|
||||
images = images.permute(0, 3, 1, 2)
|
||||
# positive_prompt_embedding 和 negative_prompt_embedding 需要匹配 batch_size
|
||||
positive_prompt_embedding = positive_prompt_embedding.repeat(batch_size, 1, 1)
|
||||
negative_prompt_embedding = negative_prompt_embedding.repeat(batch_size, 1, 1)
|
||||
|
||||
def callback(*_):
|
||||
pbar.update(1)
|
||||
|
||||
result = pipeline(
|
||||
image=images,
|
||||
generator=generator,
|
||||
prompt_embeds=positive_prompt_embedding,
|
||||
negative_prompt_embeds=negative_prompt_embedding,
|
||||
num_inference_steps=num_inference_steps,
|
||||
callback_steps=1,
|
||||
strength=strength,
|
||||
callback=callback,
|
||||
return_dict=True,
|
||||
)
|
||||
# image = result["images"][0]
|
||||
# images to torch.Tensor
|
||||
imgs = [np.array(img) for img in result["images"]]
|
||||
imgs = torch.tensor(imgs, dtype=images.dtype)
|
||||
result["images"][0].save("1.png")
|
||||
# 0 ~ 255 to 0 ~ 1
|
||||
imgs = imgs / 255
|
||||
# (B, C, H, W) to (B, H, W, C)
|
||||
return (imgs,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"GetFilledColorImage": GetFilledColorImage,
|
||||
"GetAverageColorFromImage": GetAverageColorFromImage,
|
||||
"DiffusersPipeline": DiffusersPipeline,
|
||||
"DiffusersGenerate": DiffusersGenerate,
|
||||
"DiffusersPrepareLatents": DiffusersPrepareLatents,
|
||||
"DiffusersDecoder": DiffusersDecoder,
|
||||
"DiffusersCompelPromptEmbedding": DiffusersCompelPromptEmbedding,
|
||||
"DiffusersTextureInversionLoader": DiffusersTextureInversionLoader,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GetFilledColorImage": "Get Filled Color Image Jannchie",
|
||||
"GetAverageColorFromImage": "Get Average Color From Image Jannchie",
|
||||
"DiffusersPipeline": "Diffusers Pipeline",
|
||||
"DiffusersGenerate": "Diffusers Generate",
|
||||
"DiffusersPrepareLatents": "Diffusers Prepare Latents",
|
||||
"DiffusersDecoder": "Diffusers Decoder",
|
||||
"DiffusersCompelPromptEmbedding": "Diffusers Compel Prompt Embedding",
|
||||
"DiffusersTextureInversionLoader": "Diffusers Texture Inversion Embedding Loader",
|
||||
}
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,73 @@
|
||||
import torch
|
||||
from diffusers import AutoencoderKL, StableDiffusionImg2ImgPipeline
|
||||
from diffusers.schedulers import (
|
||||
DEISMultistepScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
DPMSolverSinglestepScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
HeunDiscreteScheduler,
|
||||
KDPM2AncestralDiscreteScheduler,
|
||||
KDPM2DiscreteScheduler,
|
||||
LMSDiscreteScheduler,
|
||||
UniPCMultistepScheduler,
|
||||
)
|
||||
|
||||
import comfy.model_management
|
||||
|
||||
schedulers = {
|
||||
"DPM++ 2M": DPMSolverMultistepScheduler(),
|
||||
"DPM++ 2M Karras": DPMSolverMultistepScheduler(use_karras_sigmas=True),
|
||||
"DPM++ 2M SDE": DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++"),
|
||||
"DPM++ 2M SDE Karras": DPMSolverMultistepScheduler(
|
||||
use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
|
||||
),
|
||||
"DPM++ SDE": DPMSolverSinglestepScheduler(),
|
||||
"DPM++ SDE Karras": DPMSolverSinglestepScheduler(use_karras_sigmas=True),
|
||||
"DPM2": KDPM2DiscreteScheduler(),
|
||||
"DPM2 Karras": KDPM2DiscreteScheduler(use_karras_sigmas=True),
|
||||
"DPM2 a": KDPM2AncestralDiscreteScheduler(),
|
||||
"DPM2 a Karras": KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True),
|
||||
"Euler": EulerDiscreteScheduler(),
|
||||
"Euler a": EulerAncestralDiscreteScheduler(),
|
||||
"Heun": HeunDiscreteScheduler(),
|
||||
"LMS": LMSDiscreteScheduler(),
|
||||
"LMS Karras": LMSDiscreteScheduler(use_karras_sigmas=True),
|
||||
"DEIS": DEISMultistepScheduler(),
|
||||
"UniPC": UniPCMultistepScheduler(),
|
||||
}
|
||||
|
||||
|
||||
class PipelineWrapper:
|
||||
|
||||
def __init__(
|
||||
self, ckpt_path: str, vae_path: str = None, scheduler_name: str = None
|
||||
):
|
||||
scheduler = schedulers.get(scheduler_name)
|
||||
device = comfy.model_management.get_torch_device()
|
||||
dtype = comfy.model_management.VAE_DTYPE
|
||||
if ckpt_path.endswith(".safetensors"):
|
||||
self.pipeline = StableDiffusionImg2ImgPipeline.from_single_file(
|
||||
ckpt_path,
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
else:
|
||||
self.pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(
|
||||
ckpt_path,
|
||||
torch_dtype=dtype,
|
||||
)
|
||||
if vae_path:
|
||||
if vae_path.endswith(".safetensors"):
|
||||
self.pipeline.vae = AutoencoderKL.from_single_file(
|
||||
vae_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
else:
|
||||
self.pipeline.vae = AutoencoderKL.from_pretrained(
|
||||
vae_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
if scheduler:
|
||||
self.pipeline.scheduler = scheduler
|
||||
self.pipeline.to(device)
|
||||
self.pipeline.safety_checker = None
|
||||
Binary file not shown.
@@ -0,0 +1,3 @@
|
||||
compel
|
||||
diffusers
|
||||
numpy
|
||||
Reference in New Issue
Block a user