370 lines
15 KiB
Python
370 lines
15 KiB
Python
import torch
|
|
import os
|
|
import folder_paths
|
|
from pathlib import Path
|
|
|
|
from comfy import model_management
|
|
|
|
from .diffusers_magic_clothing.garment_diffusion import ClothAdapter
|
|
from .diffusers_magic_clothing.MagicClothingDiffusionPipeline import MagicClothingDiffusionPipeline
|
|
from diffusers import (
|
|
AutoencoderKL,
|
|
DDIMScheduler,
|
|
DDPMScheduler,
|
|
DEISMultistepScheduler,
|
|
DPMSolverMultistepScheduler,
|
|
DPMSolverSinglestepScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
EulerDiscreteScheduler,
|
|
HeunDiscreteScheduler,
|
|
KDPM2AncestralDiscreteScheduler,
|
|
KDPM2DiscreteScheduler,
|
|
UniPCMultistepScheduler,
|
|
)
|
|
|
|
SCHEDULERS = {
|
|
'DDIM' : DDIMScheduler,
|
|
'DDPM' : DDPMScheduler,
|
|
'DEISMultistep' : DEISMultistepScheduler,
|
|
'DPMSolverMultistep' : DPMSolverMultistepScheduler,
|
|
'DPMSolverSinglestep' : DPMSolverSinglestepScheduler,
|
|
'EulerAncestralDiscrete' : EulerAncestralDiscreteScheduler,
|
|
'EulerDiscrete' : EulerDiscreteScheduler,
|
|
'HeunDiscrete' : HeunDiscreteScheduler,
|
|
'KDPM2AncestralDiscrete' : KDPM2AncestralDiscreteScheduler,
|
|
'KDPM2Discrete' : KDPM2DiscreteScheduler,
|
|
'UniPCMultistep' : UniPCMultistepScheduler
|
|
}
|
|
|
|
class ChangePixelValueNormalization:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"pixels": ("IMAGE", ),
|
|
"mode": (["[0,1]=>[-1,1]", "[-1,1]=>[0,1]"],),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "normalization"
|
|
|
|
CATEGORY = "image"
|
|
|
|
def normalization(self, pixels, mode):
|
|
if mode == "[0,1]=>[-1,1]":
|
|
pixels = (pixels * 255).round().clamp(min=0, max=255) / 127.5 - 1.0
|
|
elif mode == "[-1,1]=>[0,1]":
|
|
pixels = ((pixels+1) * 127.5).clamp(min=0, max=255) / 255.0
|
|
else:
|
|
pixels = pixels
|
|
return (pixels,)
|
|
|
|
|
|
class ChangePipelineDtypeAndDevice:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"pipeline": ("PIPELINE", ),
|
|
"dtype": (["default", "float32", "float16", "bfloat16"],),
|
|
"device": (["default", "cpu", "cuda", "cuda:0", "cuda:1"],),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("PIPELINE",)
|
|
FUNCTION = "change_dtype"
|
|
|
|
CATEGORY = "pipeline"
|
|
|
|
def change_dtype(self, pipeline, dtype="default", device="default"):
|
|
if dtype == "float16":
|
|
seleted_type = torch.float16
|
|
elif dtype == "bfloat16":
|
|
seleted_type = torch.bfloat16
|
|
else:
|
|
seleted_type = torch.float32
|
|
if device == "default":
|
|
seleted_device = model_management.get_torch_device()
|
|
else:
|
|
seleted_device = torch.device(device)
|
|
pipeline = pipeline.to(seleted_device, dtype=seleted_type)
|
|
pipeline.device = seleted_device
|
|
pipeline.dtype = seleted_type
|
|
return (pipeline,)
|
|
|
|
|
|
class RunMagicClothingDiffusersModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"cloth_image": ("IMAGE",),
|
|
"magicClothingAdapter": ("MAGIC_CLOTHING_ADAPTER",),
|
|
"positive": ("STRING", {
|
|
"dynamicPrompts": False,
|
|
"multiline": True,
|
|
"default": ""
|
|
}),
|
|
"negative": ("STRING", {
|
|
"dynamicPrompts": False,
|
|
"multiline": True,
|
|
"default": ""
|
|
}),
|
|
"height": ("INT", {"default": 768, "min": 0, "max": 2048}),
|
|
"width": ("INT", {"default": 576, "min": 0, "max": 2048}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4}),
|
|
"steps": ("INT", {"default": 25, "min": 0, "max": 100}),
|
|
"cfg": ("FLOAT", {"default": 5, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"cloth_guidance_scale": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"seed": ("INT", {"default": 1234, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "run_model"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def run_model(self, cloth_image, magicClothingAdapter, positive, negative, height, width, batch_size, steps, cfg, cloth_guidance_scale, seed,):
|
|
cloth_image = (cloth_image * 255).round().clamp(min=0,
|
|
max=255).to(dtype=torch.float32) / 127.5 - 1.0
|
|
cloth_image = cloth_image.permute(0, 3, 1, 2)
|
|
if not isinstance(magicClothingAdapter, ClothAdapter):
|
|
# 如果发现不是正确的模型，就返回原始图片，不进行处理
|
|
gen_image = cloth_image.permute(0, 2, 3, 1)
|
|
gen_image = ((gen_image+1) * 127.5).clamp(min=0,
|
|
max=255).to(dtype=torch.float32) / 255.0
|
|
return (gen_image,)
|
|
magicClothingAdapter.enable_cloth_guidance = True
|
|
cloth_image = cloth_image.to(
|
|
magicClothingAdapter.pipe.device, dtype=magicClothingAdapter.pipe.dtype)
|
|
with torch.inference_mode():
|
|
prompt_embeds_null = magicClothingAdapter.pipe.encode_prompt(
|
|
[""], device=magicClothingAdapter.pipe.device, num_images_per_prompt=1, do_classifier_free_guidance=False)[0]
|
|
prompt_embeds, negative_prompt_embeds = magicClothingAdapter.pipe.encode_prompt(
|
|
positive,
|
|
magicClothingAdapter.pipe.device,
|
|
batch_size,
|
|
True,
|
|
negative,
|
|
prompt_embeds=None,
|
|
negative_prompt_embeds=None,
|
|
lora_scale=None,
|
|
clip_skip=None,
|
|
)
|
|
cloth_latent = magicClothingAdapter.pipe.vae.encode(
|
|
cloth_image).latent_dist.mode()
|
|
gen_image = magicClothingAdapter.generate(cloth_latent, None, prompt_embeds_null, prompt_embeds, negative_prompt_embeds, batch_size, seed, cfg, cloth_guidance_scale, steps, height, width)
|
|
gen_image = magicClothingAdapter.pipe.vae.decode(
|
|
gen_image, return_dict=False, generator=magicClothingAdapter.generator)[0]
|
|
gen_image = gen_image.permute(0, 2, 3, 1)
|
|
gen_image = ((gen_image+1) * 127.5).clamp(min=0,
|
|
max=255).to(dtype=torch.float32) / 255.0
|
|
return (gen_image,)
|
|
|
|
|
|
class LoadMagicClothingPipelineWithPath:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
paths = []
|
|
my_path = os.path.dirname(__file__)
|
|
my_pipeline_path = os.path.join(my_path, "conversion")
|
|
for search_path in folder_paths.get_folder_paths("diffusers"):
|
|
if os.path.exists(search_path):
|
|
client_paths = next(os.walk(search_path))[1]
|
|
client_paths = ["diffusers/" + item for item in client_paths]
|
|
paths += client_paths
|
|
if os.path.exists(my_pipeline_path):
|
|
client_paths = next(os.walk(my_pipeline_path))[1]
|
|
client_paths = ["conversion/" + item for item in client_paths]
|
|
paths += client_paths
|
|
return {"required": {"model_path": (paths,),
|
|
"dtype": (["default", "float32", "float16", "bfloat16"],),
|
|
"device": (["default", "cpu", "cuda", "cuda:0", "cuda:1"],), }}
|
|
RETURN_TYPES = ("PIPELINE", "AUTOENCODER", "SCHEDULER",)
|
|
FUNCTION = "load_checkpoint"
|
|
|
|
CATEGORY = "Diffusers"
|
|
|
|
def load_checkpoint(self, model_path,dtype,device):
|
|
if dtype == "float16":
|
|
seleted_type = torch.float16
|
|
elif dtype == "bfloat16":
|
|
seleted_type = torch.bfloat16
|
|
else:
|
|
seleted_type = torch.float32
|
|
if device == "default":
|
|
seleted_device = model_management.get_torch_device()
|
|
else:
|
|
seleted_device = torch.device(device)
|
|
|
|
if model_path.startswith("conversion/"):
|
|
model_path = model_path.replace("conversion/", "")
|
|
my_path = os.path.dirname(__file__)
|
|
my_pipeline_path = os.path.join(my_path, "conversion")
|
|
model_real_path = os.path.join(my_pipeline_path, model_path)
|
|
model_real_dir = my_pipeline_path
|
|
elif model_path.startswith("diffusers/"):
|
|
model_path = model_path.replace("diffusers/", "")
|
|
diffusers_path = folder_paths.get_folder_paths("diffusers")[0]
|
|
model_real_path = os.path.join(diffusers_path, model_path)
|
|
model_real_dir = diffusers_path
|
|
else:
|
|
raise ValueError("未选择模型")
|
|
|
|
pipe = MagicClothingDiffusionPipeline.from_pretrained(
|
|
pretrained_model_name_or_path=model_real_path,
|
|
torch_dtype=seleted_type,
|
|
cache_dir=model_real_dir,
|
|
)
|
|
pipe.to(seleted_device, dtype=seleted_type)
|
|
return ((pipe, model_real_path), pipe.vae, pipe.scheduler)
|
|
|
|
class LoadMagicClothingPipelinWithConversion:
|
|
# code base from https://github.com/Limitex/ComfyUI-Diffusers.git
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
|
|
"dtype": (["default", "float32", "float16", "bfloat16"],),
|
|
"device": (["default", "cpu", "cuda", "cuda:0", "cuda:1"],), }}
|
|
|
|
RETURN_TYPES = ("PIPELINE", "AUTOENCODER", "SCHEDULER",)
|
|
|
|
FUNCTION = "create_pipeline"
|
|
|
|
CATEGORY = "Diffusers"
|
|
|
|
def create_pipeline(self, ckpt_name,dtype,device):
|
|
if dtype == "float16":
|
|
seleted_type = torch.float16
|
|
elif dtype == "bfloat16":
|
|
seleted_type = torch.bfloat16
|
|
else:
|
|
seleted_type = torch.float32
|
|
if device == "default":
|
|
seleted_device = model_management.get_torch_device()
|
|
else:
|
|
seleted_device = torch.device(device)
|
|
my_path = os.path.dirname(__file__)
|
|
my_pipeline_path = os.path.join(my_path, "conversion")
|
|
if not os.path.exists(my_pipeline_path):
|
|
os.makedirs(my_pipeline_path)
|
|
real_ckpt_name = Path(ckpt_name).stem
|
|
real_ckpt_name = real_ckpt_name +"_"+str(seleted_type)
|
|
real_ckpt_name = real_ckpt_name.replace(" ", "_").replace(".", "_").replace("/", "_")
|
|
ckpt_conversion_path = os.path.join(my_pipeline_path, real_ckpt_name)
|
|
if not os.path.exists(ckpt_conversion_path):
|
|
# 不存在，则进行转换
|
|
MagicClothingDiffusionPipeline.from_single_file(
|
|
pretrained_model_link_or_path=folder_paths.get_full_path("checkpoints", ckpt_name),
|
|
torch_dtype=seleted_type,
|
|
cache_dir=my_pipeline_path,
|
|
).save_pretrained(ckpt_conversion_path, safe_serialization=True)
|
|
|
|
pipe = MagicClothingDiffusionPipeline.from_pretrained(
|
|
pretrained_model_name_or_path=ckpt_conversion_path,
|
|
torch_dtype=seleted_type,
|
|
cache_dir=my_pipeline_path,
|
|
)
|
|
pipe.to(seleted_device, dtype=seleted_type)
|
|
return ((pipe, ckpt_conversion_path), pipe.vae, pipe.scheduler)
|
|
|
|
|
|
|
|
class DiffusersSchedulerLoader:
|
|
# code copy from https://github.com/Limitex/ComfyUI-Diffusers.git
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("PIPELINE", ),
|
|
"scheduler_name": (list(SCHEDULERS.keys()), ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("SCHEDULER",)
|
|
|
|
FUNCTION = "load_scheduler"
|
|
|
|
CATEGORY = "Diffusers"
|
|
|
|
def load_scheduler(self, pipeline, scheduler_name):
|
|
my_path = os.path.dirname(__file__)
|
|
my_pipeline_path = os.path.join(my_path, "conversion")
|
|
if not os.path.exists(my_pipeline_path):
|
|
os.makedirs(my_pipeline_path)
|
|
scheduler = SCHEDULERS[scheduler_name].from_pretrained(
|
|
pretrained_model_name_or_path=pipeline[1],
|
|
torch_dtype=pipeline[0].dtype,
|
|
cache_dir=my_pipeline_path,
|
|
subfolder='scheduler'
|
|
)
|
|
return (scheduler,)
|
|
|
|
class DiffusersModelMakeup:
|
|
# code copy from https://github.com/Limitex/ComfyUI-Diffusers.git
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("PIPELINE", ),
|
|
"scheduler": ("SCHEDULER", ),
|
|
"autoencoder": ("AUTOENCODER", ),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MAKED_PIPELINE",)
|
|
|
|
FUNCTION = "makeup_pipeline"
|
|
|
|
CATEGORY = "Diffusers"
|
|
|
|
def makeup_pipeline(self, pipeline, scheduler, autoencoder):
|
|
pipeline = pipeline[0]
|
|
autoencoder.to(pipeline.device, dtype=pipeline.dtype)
|
|
pipeline.vae = autoencoder
|
|
pipeline.scheduler = scheduler
|
|
pipeline.safety_checker = None if pipeline.safety_checker is None else lambda images, **kwargs: (images, [False])
|
|
pipeline.enable_attention_slicing()
|
|
return (pipeline,)
|
|
|
|
class LoadMagicClothingAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"magicClothingUnet": (folder_paths.get_filename_list("unet"), ),
|
|
"pipeline": ("MAKED_PIPELINE", ),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MAGIC_CLOTHING_ADAPTER",)
|
|
RETURN_NAMES = ("MagicClothingAdapter",)
|
|
FUNCTION = "load_model"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_model(self, magicClothingUnet, pipeline):
|
|
unet_path = folder_paths.get_full_path("unet", magicClothingUnet)
|
|
full_model = ClothAdapter(pipeline, unet_path)
|
|
return (full_model,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Diffusers Model Makeup &MC": DiffusersModelMakeup,
|
|
"Diffusers Scheduler Loader &MC": DiffusersSchedulerLoader,
|
|
"Change Pixel Value Normalization": ChangePixelValueNormalization,
|
|
"Change Pipeline Dtype And Device": ChangePipelineDtypeAndDevice,
|
|
"Load Magic Clothing Pipeline With Path": LoadMagicClothingPipelineWithPath,
|
|
"Load Magic Clothing Pipeline": LoadMagicClothingPipelinWithConversion,
|
|
"Load Magic Clothing Adapter": LoadMagicClothingAdapter,
|
|
"RUN Magic Clothing Diffusers Model": RunMagicClothingDiffusersModel,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Diffusers Model Makeup &MC": "Diffusers Model Makeup &MC",
|
|
"Diffusers Scheduler Loader &MC": "Diffusers Scheduler Loader &MC",
|
|
"Change Pipeline Dtype And Device": "Change Pipeline Dtype And Device",
|
|
"Change Pixel Value Normalization": "Change Pixel Value Normalization",
|
|
"Load Magic Clothing Pipeline With Path":"Load Magic Clothing Pipeline With Path&Diffusers",
|
|
"Load Magic Clothing Pipeline":"Load Magic Clothing Pipeline&Diffusers",
|
|
"Load Magic Clothing Adapter": "Load Magic Clothing Adapter &Diffusers",
|
|
"RUN Magic Clothing Adapter": "RUN Magic Clothing Adapter &Diffusers",
|
|
}
|