# 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 diffusers import diffusers.image_processor as diffusers_image_processor import torch import transformers from . import util from .util import DEVICES, DTYPES 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 = util.mkkwargs(kwargs) to_device = util.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() vae = None if hasattr(pipeline, "vae"): vae = pipeline.vae return (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 = util.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 = util.mkkwargs(kwargs) to_device = util.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 HFDRunPipeline: """ Run a HuggingFace Diffusers pipeline. """ @classmethod def INPUT_TYPES(cls): return { "required": { "pipeline": ("HFD_PIPELINE",), "prompt": ("STRING", {"default": "a photo of an astronaut riding a horse on mars", "multiline": True}), "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, width, height, seed, num_inference_steps, **kwargs ): kwargs = util.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 width > 0: kwargs["width"] = width if height > 0: kwargs["height"] = height 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 hasattr(r, "frames"): # Video model if kwargs["output_type"] == "pil": return (r.frames[0], None) else: return (None, r.frames) else: # Image model 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", "multiline": True}), "kwargs": ("STRING",), } } RETURN_TYPES = ("TENSOR", "TENSOR", "TENSOR", "TENSOR") FUNCTION = "encode" CATEGORY = "huggingface-diffusers" def encode(self, pipeline, prompt, kwargs): kwargs = util.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,)