151 lines
4.5 KiB
Python
151 lines
4.5 KiB
Python
import os
|
|
from enum import Enum
|
|
import torch
|
|
|
|
import folder_paths
|
|
import comfy.model_management
|
|
from comfy.model_patcher import ModelPatcher
|
|
from comfy.utils import load_torch_file
|
|
from .lib_layerdiffusion.utils import (
|
|
rgba2rgbfp32,
|
|
load_file_from_url,
|
|
to_lora_patch_dict,
|
|
)
|
|
from .lib_layerdiffusion.models import TransparentVAEDecoder
|
|
|
|
|
|
layer_model_root = os.path.join(folder_paths.models_dir, "layer_model")
|
|
load_layer_model_state_dict = load_torch_file
|
|
|
|
|
|
class RGBA2RBGfp32:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "rgb2rgb_fp32"
|
|
CATEGORY = "layered_diffusion"
|
|
|
|
def rgb2rgb_fp32(self, image):
|
|
return rgba2rgbfp32(image)
|
|
|
|
|
|
class LayeredDiffusionDecode:
|
|
"""
|
|
Decode alpha channel value from pixel value.
|
|
[B, C=3, H, W] => [B, C=4, H, W]
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"samples": ("LATENT",), "images": ("IMAGE",)}}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "decode"
|
|
CATEGORY = "layered_diffusion"
|
|
|
|
def __init__(self) -> None:
|
|
self.vae_transparent_decoder = None
|
|
|
|
def decode(self, samples, images):
|
|
if self.vae_transparent_decoder is None:
|
|
model_path = load_file_from_url(
|
|
url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_decoder.safetensors",
|
|
model_dir=layer_model_root,
|
|
file_name="vae_transparent_decoder.safetensors",
|
|
)
|
|
self.vae_transparent_decoder = TransparentVAEDecoder(
|
|
load_torch_file(model_path),
|
|
device=comfy.model_management.get_torch_device(),
|
|
dtype=(
|
|
torch.float16
|
|
if comfy.model_management.should_use_fp16()
|
|
else torch.float32
|
|
),
|
|
)
|
|
latent = samples["samples"]
|
|
pixel = images.movedim(-1, 1) # [B, H, W, C] => [B, C, H, W]
|
|
pixel_with_alpha = self.vae_transparent_decoder.decode_pixel(pixel, latent)
|
|
# [B, C, H, W] => [B, H, W, C]
|
|
pixel_with_alpha = pixel_with_alpha.movedim(1, -1)
|
|
image = pixel_with_alpha[..., 1:]
|
|
alpha = pixel_with_alpha[..., 0]
|
|
return (image, alpha)
|
|
|
|
|
|
class LayerMethod(Enum):
|
|
ATTN = "Attention Injection"
|
|
CONV = "Conv Injection"
|
|
|
|
|
|
class LayeredDiffusionApply:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"method": (
|
|
[
|
|
LayerMethod.ATTN.value,
|
|
LayerMethod.CONV.value,
|
|
],
|
|
{
|
|
"default": LayerMethod.ATTN.value,
|
|
},
|
|
),
|
|
"weight": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": -1, "max": 3, "step": 0.05},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "apply_layered_diffusion"
|
|
CATEGORY = "layered_diffusion"
|
|
|
|
def apply_layered_diffusion(
|
|
self,
|
|
model: ModelPatcher,
|
|
method: str,
|
|
weight: float,
|
|
):
|
|
"""Patch model"""
|
|
method = LayerMethod(method)
|
|
|
|
# Patch unet
|
|
if method == LayerMethod.ATTN:
|
|
model_path = load_file_from_url(
|
|
url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_attn.safetensors",
|
|
model_dir=layer_model_root,
|
|
file_name="layer_xl_transparent_attn.safetensors",
|
|
)
|
|
if method == LayerMethod.CONV:
|
|
model_path = load_file_from_url(
|
|
url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/layer_xl_transparent_conv.safetensors",
|
|
model_dir=layer_model_root,
|
|
file_name="layer_xl_transparent_conv.safetensors",
|
|
)
|
|
|
|
layer_lora_state_dict = load_layer_model_state_dict(model_path)
|
|
layer_lora_patch_dict = to_lora_patch_dict(layer_lora_state_dict)
|
|
work_model = model.clone()
|
|
work_model.add_patches(layer_lora_patch_dict, weight)
|
|
return (work_model,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LayeredDiffusionApply": LayeredDiffusionApply,
|
|
"LayeredDiffusionDecode": LayeredDiffusionDecode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LayeredDiffusionApply": "Layer Diffusion Apply",
|
|
"LayeredDiffusionDecode": "Layer Diffusion Decode",
|
|
}
|