388 lines
12 KiB
Python
388 lines
12 KiB
Python
# This is free and unencumbered software released into the public domain.
|
|
#
|
|
# Anyone is free to copy, modify, publish, use, compile, sell, or distribute
|
|
# this software, either in source code form or as a compiled binary, for any
|
|
# purpose, commercial or non-commercial, and by any means.
|
|
#
|
|
# In jurisdictions that recognize copyright laws, the author or authors of this
|
|
# software dedicate any and all copyright interest in the software to the public
|
|
# domain. We make this dedication for the benefit of the public at large and to
|
|
# the detriment of our heirs and successors. We intend this dedication to be an
|
|
# overt act of relinquishment in perpetuity of all present and future rights to
|
|
# this software under copyright law.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
|
|
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
|
|
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
|
|
import json
|
|
|
|
import diffusers
|
|
import diffusers.image_processor as diffusers_image_processor
|
|
import torch
|
|
import transformers
|
|
|
|
DEVICES = ["default", "auto", "cpu"]
|
|
DEFAULT_DEVICE = "cpu"
|
|
if torch.cuda.is_available():
|
|
DEVICES.append("cuda")
|
|
for i in range(torch.cuda.device_count()):
|
|
DEVICES.append(f"cuda:{i}")
|
|
DEFAULT_DEVICE = "cuda"
|
|
DTYPES = ("default", "float32", "bfloat16", "float16", "bitsandbytes_8bit", "bitsandbytes_4bit")
|
|
|
|
def get_device(device):
|
|
if device == "default":
|
|
return DEFAULT_DEVICE
|
|
else:
|
|
return device
|
|
|
|
def mkkwargs(kwargs):
|
|
"""
|
|
Return kwargs appropriate for HuggingFace models.
|
|
"""
|
|
ret = {}
|
|
if "kwargs" in kwargs and kwargs["kwargs"] != "":
|
|
ret = json.loads(kwargs["kwargs"])
|
|
for key in kwargs:
|
|
if key == "kwargs":
|
|
continue
|
|
if key not in ret and kwargs[key] is not None:
|
|
ret[key] = kwargs[key]
|
|
return ret
|
|
|
|
def apply_device(kwargs, device, dtype, enable_model_cpu_offload=False, quant="pipeline"):
|
|
"""
|
|
Apply device and dtype properties to the kwargs. Quantizes in pipeline,
|
|
transformers, or diffusers mode. Returns the device that the result should
|
|
be moved to with `to`, or `None` if not needed.
|
|
"""
|
|
device = get_device(device)
|
|
to_device = None
|
|
if not enable_model_cpu_offload:
|
|
if ":" in device:
|
|
to_device = device
|
|
else:
|
|
kwargs["device_map"] = get_device(device)
|
|
kwargs["torch_dtype"] = torch.bfloat16
|
|
|
|
if dtype[0:13] == "bitsandbytes_":
|
|
if quant == "transformers" or quant == "diffusers":
|
|
qc = {}
|
|
if dtype == "bitsandbytes_4bit":
|
|
qc["load_in_4bit"] = True
|
|
else:
|
|
qc["load_in_8bit"] = True
|
|
if quant == "transformers":
|
|
qc = transformers.BitsAndBytesConfig(**qc)
|
|
else:
|
|
qc = diffusers.BitsAndBytesConfig(**qc)
|
|
kwargs["quantization_config"] = qc
|
|
|
|
else: # "pipeline"
|
|
qc = {
|
|
"quant_backend": dtype
|
|
}
|
|
if dtype == "bitsandbytes_4bit":
|
|
qc["quant_kwargs"] = {"load_in_4bit": True}
|
|
else:
|
|
qc["quant_kwargs"] = {"load_in_8bit": True}
|
|
kwargs["quantization_config"] = diffusers.PipelineQuantizationConfig(**qc)
|
|
|
|
elif dtype == "float32":
|
|
kwargs["torch_dtype"] = torch.float32
|
|
elif dtype == "float16":
|
|
kwargs["torch_dtype"] = torch.float16
|
|
|
|
return to_device
|
|
|
|
class HFDLoadPipeline:
|
|
"""
|
|
Load a HuggingFace Diffusers pipeline.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline_class": ("STRING", {"default": "AutoPipelineForText2Image"}),
|
|
"model": ("STRING", {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
|
|
"device": (DEVICES,),
|
|
"enable_model_cpu_offload": ([False, True],),
|
|
"dtype": (DTYPES,),
|
|
"kwargs": ("STRING",),
|
|
},
|
|
"optional": {
|
|
"vae": ("HFD_AUTOENCODERKL",),
|
|
"text_encoder": ("HFT_MODEL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HFD_PIPELINE", "HFD_AUTOENCODERKL")
|
|
FUNCTION = "load"
|
|
|
|
CATEGORY = "huggingface-diffusers"
|
|
|
|
def load(
|
|
self, pipeline_class, model, device, enable_model_cpu_offload, dtype,
|
|
**kwargs
|
|
):
|
|
kwargs = mkkwargs(kwargs)
|
|
to_device = apply_device(
|
|
kwargs, device, dtype,
|
|
enable_model_cpu_offload=enable_model_cpu_offload
|
|
)
|
|
pipeline = getattr(diffusers, pipeline_class).from_pretrained(
|
|
model, **kwargs
|
|
)
|
|
if to_device:
|
|
pipeline.to(to_device)
|
|
if enable_model_cpu_offload:
|
|
pipeline.enable_model_cpu_offload()
|
|
return (pipeline, pipeline.vae)
|
|
|
|
class HFDLoadLora:
|
|
"""
|
|
Load a Lora into a HuggingFace pipeline.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HFD_PIPELINE",),
|
|
"pretrained_model_name_or_path_or_dict": ("STRING", {"default": "TheLastBen/Papercut_SDXL"}),
|
|
"weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}),
|
|
"kwargs": ("STRING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HFD_PIPELINE",)
|
|
FUNCTION = "load"
|
|
|
|
CATEGORY = "huggingface-diffusers"
|
|
|
|
def load(self, pipeline, weight, **kwargs):
|
|
kwargs = mkkwargs(kwargs)
|
|
pipeline.load_lora_weights(**kwargs)
|
|
pipeline.fuse_lora(lora_scale=weight)
|
|
pipeline.unload_lora_weights()
|
|
return (pipeline,)
|
|
|
|
class HFDAutoencoderKL:
|
|
"""
|
|
Load an Autoencoder.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"autoencoder_class": ("STRING", {"default": "AutoencoderKL"}),
|
|
"pretrained_model_name_or_path": ("STRING", {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
|
|
"subfolder": ("STRING", {"default": "vae"}),
|
|
"device": (DEVICES,),
|
|
"dtype": (DTYPES,),
|
|
"kwargs": ("STRING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HFD_AUTOENCODERKL",)
|
|
FUNCTION = "load"
|
|
|
|
CATEGORY = "huggingface-diffusers"
|
|
|
|
def load(self, autoencoder_class, device, dtype, **kwargs):
|
|
kwargs = mkkwargs(kwargs)
|
|
to_device = apply_device(kwargs, device, dtype, quant="diffusers")
|
|
vae = getattr(diffusers, autoencoder_class).from_pretrained(**kwargs)
|
|
if to_device:
|
|
vae.to(to_device)
|
|
return (vae,)
|
|
|
|
class HFTAutoModel:
|
|
"""
|
|
Load a transformers model.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_class": ("STRING", {"default": "AutoModel"}),
|
|
"pretrained_model_name_or_path": ("STRING", {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
|
|
"subfolder": ("STRING", {"default": "text_encoder"}),
|
|
"device": (DEVICES,),
|
|
"dtype": (DTYPES,),
|
|
"kwargs": ("STRING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("HFT_MODEL",)
|
|
FUNCTION = "load"
|
|
|
|
CATEGORY = "huggingface-transformers"
|
|
|
|
def load(self, model_class, device, dtype, **kwargs):
|
|
kwargs = mkkwargs(kwargs)
|
|
to_device = apply_device(kwargs, device, dtype, quant="transformers")
|
|
model = getattr(transformers, model_class).from_pretrained(**kwargs)
|
|
if to_device:
|
|
model.to(to_device)
|
|
return (model,)
|
|
|
|
class HFDRunPipeline:
|
|
"""
|
|
Run HuggingFace pipelines.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HFD_PIPELINE",),
|
|
"prompt": ("STRING", {"default": "a photo of an astronaut riding a horse on mars"}),
|
|
"negative_prompt": ("STRING",),
|
|
"width": ("INT", {"default": 1024}),
|
|
"height": ("INT", {"default": 1024}),
|
|
"seed": ("INT", {"default": 1, "min": 0, "max": 0xffffffffffffffff}),
|
|
"num_inference_steps": ("INT", {"default": 0, "min": 0, "max": 0x10000}),
|
|
"output_type": (["pil", "latent"],),
|
|
"kwargs": ("STRING",),
|
|
},
|
|
"optional": {
|
|
"image": ("PIL_IMAGE",),
|
|
"mask_image": ("PIL_IMAGE",),
|
|
"latents": ("LATENT",),
|
|
"prompt_embeds": ("TENSOR",),
|
|
"pooled_prompt_embeds": ("TENSOR",),
|
|
"negative_prompt_embeds": ("TENSOR",),
|
|
"negative_pooled_prompt_embeds": ("TENSOR",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("PIL_IMAGE", "LATENT")
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "huggingface-diffusers"
|
|
|
|
def generate(
|
|
self, negative_prompt, pipeline, seed, num_inference_steps, **kwargs
|
|
):
|
|
kwargs = mkkwargs(kwargs)
|
|
if "prompt_embeds" in kwargs:
|
|
del kwargs["prompt"]
|
|
if "negative_prompt_embeds" not in kwargs and negative_prompt != "":
|
|
kwargs["negative_prompt"] = negative_prompt
|
|
if num_inference_steps > 0:
|
|
kwargs["num_inference_steps"] = num_inference_steps
|
|
with torch.no_grad():
|
|
generator = torch.manual_seed(seed)
|
|
r = pipeline(
|
|
generator=generator,
|
|
**kwargs
|
|
)
|
|
if kwargs["output_type"] == "pil":
|
|
return (r.images[0], None)
|
|
else:
|
|
return (None, r.images)
|
|
|
|
class HFDEncodePrompt:
|
|
"""
|
|
Encode a prompt using a HuggingFace pipeline.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("HFD_PIPELINE",),
|
|
"prompt": ("STRING", {"default": "a photo of an astronaut riding a horse on mars"}),
|
|
"kwargs": ("STRING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TENSOR", "TENSOR", "TENSOR", "TENSOR")
|
|
FUNCTION = "encode"
|
|
|
|
CATEGORY = "huggingface-diffusers"
|
|
|
|
def encode(self, pipeline, prompt, kwargs):
|
|
kwargs = mkkwargs(kwargs)
|
|
r = list(pipeline.encode_prompt(prompt, **kwargs))
|
|
while len(r) < 4:
|
|
r.append(None)
|
|
return tuple(r)
|
|
|
|
class HFDVAEDecode:
|
|
"""
|
|
VAE decoding using HuggingFace diffusers.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"latents": ("LATENT",),
|
|
"vae": ("HFD_AUTOENCODERKL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("PIL_IMAGE",)
|
|
FUNCTION = "decode"
|
|
|
|
CATEGORY = "huggingface-diffusers"
|
|
|
|
def decode(self, latents, vae):
|
|
image_processor = diffusers_image_processor.VaeImageProcessor(
|
|
vae_scale_factor=2**(len(vae.config.block_out_channels) - 1)
|
|
)
|
|
with torch.no_grad():
|
|
image = vae.decode(
|
|
(latents / vae.config.scaling_factor).to(device=vae.device, dtype=vae.dtype),
|
|
return_dict=False
|
|
)[0]
|
|
image = image_processor.postprocess(image)
|
|
return (image,)
|
|
|
|
class HFDVAEEncode:
|
|
"""
|
|
VAE encoding using HuggingFace diffusers.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("PIL_IMAGE",),
|
|
"vae": ("HFD_AUTOENCODERKL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "encode"
|
|
|
|
CATEGORY = "huggingface-diffusers"
|
|
|
|
def encode(self, image, vae):
|
|
image_processor = diffusers_image_processor.VaeImageProcessor(
|
|
vae_scale_factor=2**(len(vae.config.block_out_channels) - 1)
|
|
)
|
|
if isinstance(image, list):
|
|
image = image[0]
|
|
latents = image_processor.preprocess(
|
|
image,
|
|
height=image.height,
|
|
width=image.width
|
|
).to(device=vae.device, dtype=vae.dtype)
|
|
with torch.no_grad():
|
|
latents = vae.encode(
|
|
latents,
|
|
return_dict=False
|
|
)[0].sample()
|
|
latents *= vae.config.scaling_factor
|
|
return (latents,)
|