Initial import
This commit is contained in:
Yahweasel
@@ -0,0 +1,387 @@
|
||||
# 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,)
|
||||
Reference in New Issue
Block a user