|
|
import folder_paths
|
|
|
import os
|
|
|
import torch
|
|
|
import torch.nn.functional as F
|
|
|
from comfy.utils import ProgressBar, load_torch_file
|
|
|
import comfy.sample
|
|
|
from nodes import CLIPTextEncode
|
|
|
|
|
|
script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
folder_paths.add_model_folder_path("intrinsic_loras", os.path.join(script_directory, "intrinsic_loras"))
|
|
|
|
|
|
class Intrinsic_lora_sampling:
|
|
|
def __init__(self):
|
|
|
self.loaded_lora = None
|
|
|
|
|
|
@classmethod
|
|
|
def INPUT_TYPES(s):
|
|
|
return {"required": { "model": ("MODEL",),
|
|
|
"lora_name": (folder_paths.get_filename_list("intrinsic_loras"), ),
|
|
|
"task": (
|
|
|
[
|
|
|
'depth map',
|
|
|
'surface normals',
|
|
|
'albedo',
|
|
|
'shading',
|
|
|
],
|
|
|
{
|
|
|
"default": 'depth map'
|
|
|
}),
|
|
|
"text": ("STRING", {"multiline": True, "default": ""}),
|
|
|
"clip": ("CLIP", ),
|
|
|
"vae": ("VAE", ),
|
|
|
"per_batch": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}),
|
|
|
},
|
|
|
"optional": {
|
|
|
"image": ("IMAGE",),
|
|
|
"optional_latent": ("LATENT",),
|
|
|
},
|
|
|
}
|
|
|
|
|
|
RETURN_TYPES = ("IMAGE", "LATENT",)
|
|
|
FUNCTION = "onestepsample"
|
|
|
CATEGORY = "KJNodes"
|
|
|
DESCRIPTION = """
|
|
|
Sampler to use the intrinsic loras:
|
|
|
https://github.com/duxiaodan/intrinsic-lora
|
|
|
These LoRAs are tiny and thus included
|
|
|
with this node pack.
|
|
|
"""
|
|
|
|
|
|
def onestepsample(self, model, lora_name, clip, vae, text, task, per_batch, image=None, optional_latent=None):
|
|
|
pbar = ProgressBar(3)
|
|
|
|
|
|
if optional_latent is None:
|
|
|
image_list = []
|
|
|
for start_idx in range(0, image.shape[0], per_batch):
|
|
|
sub_pixels = vae.vae_encode_crop_pixels(image[start_idx:start_idx+per_batch])
|
|
|
image_list.append(vae.encode(sub_pixels[:,:,:,:3]))
|
|
|
sample = torch.cat(image_list, dim=0)
|
|
|
else:
|
|
|
sample = optional_latent["samples"]
|
|
|
noise = torch.zeros(sample.size(), dtype=sample.dtype, layout=sample.layout, device="cpu")
|
|
|
prompt = task + "," + text
|
|
|
positive, = CLIPTextEncode.encode(self, clip, prompt)
|
|
|
negative = positive
|
|
|
|
|
|
pbar.update(1)
|
|
|
|
|
|
|
|
|
class X0_PassThrough(comfy.model_sampling.EPS):
|
|
|
def calculate_denoised(self, sigma, model_output, model_input):
|
|
|
return model_output
|
|
|
def calculate_input(self, sigma, noise):
|
|
|
return noise
|
|
|
sampling_base = comfy.model_sampling.ModelSamplingDiscrete
|
|
|
sampling_type = X0_PassThrough
|
|
|
|
|
|
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
|
|
pass
|
|
|
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
|
|
|
|
|
|
|
|
model_clone = model.clone()
|
|
|
lora_path = folder_paths.get_full_path("intrinsic_loras", lora_name)
|
|
|
lora = load_torch_file(lora_path, safe_load=True)
|
|
|
self.loaded_lora = (lora_path, lora)
|
|
|
|
|
|
model_clone_with_lora = comfy.sd.load_lora_for_models(model_clone, None, lora, 1.0, 0)[0]
|
|
|
|
|
|
model_clone_with_lora.add_object_patch("model_sampling", model_sampling)
|
|
|
|
|
|
samples = {"samples": comfy.sample.sample(model_clone_with_lora, noise, 1, 1.0, "euler", "simple", positive, negative, sample,
|
|
|
denoise=1.0, disable_noise=True, start_step=0, last_step=1,
|
|
|
force_full_denoise=True, noise_mask=None, callback=None, disable_pbar=True, seed=None)}
|
|
|
pbar.update(1)
|
|
|
|
|
|
decoded = []
|
|
|
for start_idx in range(0, samples["samples"].shape[0], per_batch):
|
|
|
decoded.append(vae.decode(samples["samples"][start_idx:start_idx+per_batch]))
|
|
|
image_out = torch.cat(decoded, dim=0)
|
|
|
|
|
|
pbar.update(1)
|
|
|
|
|
|
if task == 'depth map':
|
|
|
imax = image_out.max()
|
|
|
imin = image_out.min()
|
|
|
image_out = (image_out-imin)/(imax-imin)
|
|
|
image_out = torch.max(image_out, dim=3, keepdim=True)[0].repeat(1, 1, 1, 3)
|
|
|
elif task == 'surface normals':
|
|
|
image_out = F.normalize(image_out * 2 - 1, dim=3) / 2 + 0.5
|
|
|
image_out = 1.0 - image_out
|
|
|
else:
|
|
|
image_out = image_out.clamp(-1.,1.)
|
|
|
|
|
|
return (image_out, samples,) |
