506 lines
16 KiB
Python
506 lines
16 KiB
Python
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",
|
|
}
|