338 lines
14 KiB
Python
338 lines
14 KiB
Python
import torch
|
|
import gc
|
|
import os
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from folder_paths import map_legacy, folder_names_and_paths
|
|
from .controlnet_union import ControlNetModel_Union
|
|
from .pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
|
|
from diffusers import AutoencoderKL, TCDScheduler
|
|
from diffusers.models.model_loading_utils import load_state_dict
|
|
|
|
|
|
# Get the absolute path of various directories
|
|
my_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
class PadImageForDiffusersOutpaint:
|
|
_alignment_options = ["Middle", "Left", "Right", "Top", "Bottom"]
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"width": ("INT", {"default": 720, "min": 320, "max": 1536, "tooltip": "The width used for the image."}),
|
|
"height": ("INT", {"default": 1280, "min": 320, "max": 1536, "tooltip": "The height used for the image."}),
|
|
"alignment": (s._alignment_options, {"tooltip": "Where the original image should be in the outpainted one"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
|
|
RETURN_NAMES = ("IMAGE", "MASK", "diffuser_outpaint_cnet_image")
|
|
FUNCTION = "expand_image"
|
|
CATEGORY = "DiffusersOutpaint"
|
|
|
|
def expand_image(self, image, width, height, alignment="Middle"):
|
|
|
|
# Resize Image
|
|
def can_expand(source_width, source_height, target_width, target_height, alignment):
|
|
"""Checks if the image can be expanded based on the alignment."""
|
|
if alignment in ("Left", "Right") and source_width >= target_width:
|
|
return False
|
|
if alignment in ("Top", "Bottom") and source_height >= target_height:
|
|
return False
|
|
return True
|
|
|
|
im=tensor2pil(image)
|
|
source=im.convert('RGB')
|
|
target_size = (width, height)
|
|
|
|
# Upscale if source is smaller than target in both dimensions
|
|
if source.width < target_size[0] and source.height < target_size[1]:
|
|
scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)
|
|
new_width = int(source.width * scale_factor)
|
|
new_height = int(source.height * scale_factor)
|
|
source = source.resize((new_width, new_height), Image.LANCZOS)
|
|
|
|
if source.width > target_size[0] or source.height > target_size[1]:
|
|
scale_factor = min(target_size[0] / source.width, target_size[1] / source.height)
|
|
new_width = int(source.width * scale_factor)
|
|
new_height = int(source.height * scale_factor)
|
|
source = source.resize((new_width, new_height), Image.LANCZOS)
|
|
|
|
if not can_expand(source.width, source.height, target_size[0], target_size[1], alignment):
|
|
alignment = "Middle"
|
|
# Calculate margins based on alignment
|
|
if alignment == "Middle":
|
|
margin_x = (target_size[0] - source.width) // 2
|
|
margin_y = (target_size[1] - source.height) // 2
|
|
elif alignment == "Left":
|
|
margin_x = 0
|
|
margin_y = (target_size[1] - source.height) // 2
|
|
elif alignment == "Right":
|
|
margin_x = target_size[0] - source.width
|
|
margin_y = (target_size[1] - source.height) // 2
|
|
elif alignment == "Top":
|
|
margin_x = (target_size[0] - source.width) // 2
|
|
margin_y = 0
|
|
elif alignment == "Bottom":
|
|
margin_x = (target_size[0] - source.width) // 2
|
|
margin_y = target_size[1] - source.height
|
|
|
|
background = Image.new('RGB', target_size, (255, 255, 255))
|
|
background.paste(source, (margin_x, margin_y))
|
|
|
|
image=pil2tensor(background)
|
|
#----------------------------------------------------
|
|
d1, d2, d3, d4 = image.size()
|
|
left, top, bottom, right = 0, 0, 0, 0
|
|
# Image
|
|
new_image = torch.ones(
|
|
(d1, d2 + top + bottom, d3 + left + right, d4),
|
|
dtype=torch.float32,
|
|
) * 0.5
|
|
new_image[:, top:top + d2, left:left + d3, :] = image
|
|
#----------------------------------------------------
|
|
# Mask coordinates
|
|
if alignment == "Middle":
|
|
margin_x = (width - new_width) // 2
|
|
margin_y = (height - new_height) // 2
|
|
elif alignment == "Left":
|
|
margin_x = 0
|
|
margin_y = (height - new_height) // 2
|
|
elif alignment == "Right":
|
|
margin_x = width - new_width
|
|
margin_y = (height - new_height) // 2
|
|
elif alignment == "Top":
|
|
margin_x = (width - new_width) // 2
|
|
margin_y = 0
|
|
elif alignment == "Bottom":
|
|
margin_x = (width - new_width) // 2
|
|
margin_y = height - new_height
|
|
|
|
# Create mask as big as new img
|
|
mask = torch.ones(
|
|
(height, width),
|
|
dtype=torch.float32,
|
|
)
|
|
# Create hole in mask
|
|
t = torch.zeros(
|
|
(new_height, new_width),
|
|
dtype=torch.float32
|
|
)
|
|
# Create holed mask
|
|
mask[margin_y:margin_y + new_height,
|
|
margin_x:margin_x + new_width
|
|
] = t
|
|
#----------------------------------------------------
|
|
# Prepare "cn_image" for diffusers outpaint
|
|
im=tensor2pil(new_image)
|
|
pil_new_image=im.convert('RGB')
|
|
|
|
pil_mask=tensor2pil(mask)
|
|
|
|
cnet_image = pil_new_image.copy() # copy background as cnet_image
|
|
cnet_image.paste(0, (0, 0), pil_mask) # paste mask over cnet_image, cropping it a bit
|
|
|
|
tensor_cnet_image=pil2tensor(cnet_image)
|
|
|
|
return (new_image, mask, tensor_cnet_image,)
|
|
|
|
|
|
def get_first_folder_list(folder_name: str) -> tuple[list[str], dict[str, float], float]:
|
|
folder_name = map_legacy(folder_name)
|
|
global folder_names_and_paths
|
|
folders = folder_names_and_paths[folder_name]
|
|
if folder_name == "unet":
|
|
root_folder = folders[0][0]
|
|
elif folder_name == "diffusion_models":
|
|
root_folder = folders[0][1]
|
|
visible_folders = [name for name in os.listdir(root_folder) if os.path.isdir(os.path.join(root_folder, name))]
|
|
return visible_folders
|
|
|
|
|
|
class LoadDiffusersOutpaintModels:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (get_first_folder_list("diffusion_models"), {"default": "RealVisXL_V5.0_Lightning", "tooltip": "The diffuser model used for denoising the input latent. (Put model files in a folder, in diffusion_models folder)."}),
|
|
"vae": (get_first_folder_list("diffusion_models"), {"default": "sdxl-vae-fp16-fix", "tooltip": "The vae model used for denoising the input latent. (Put model files in a folder, in diffusion_models folder)."}),
|
|
"controlnet_model": (get_first_folder_list("diffusion_models"), {"default": "controlnet-union-sdxl-1.0", "tooltip": "The controlnet model used for denoising the input latent. (Put model files in a folder, in diffusion_models folder)."}),
|
|
},
|
|
"optional": {
|
|
"enable_vae_slicing": ("BOOLEAN", {"default": True, "tooltip": "VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes."}),
|
|
"enable_vae_tiling": ("BOOLEAN", {"default": False, "tooltip": "Drastically reduces memory use but may introduce seams"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("PIPE",)
|
|
RETURN_NAMES = ("diffusers_outpaint_pipe",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "DiffusersOutpaint"
|
|
|
|
def load(self, model, vae, controlnet_model, enable_vae_slicing, enable_vae_tiling):
|
|
# Go 2 folders back
|
|
comfy_dir = os.path.dirname(os.path.dirname(my_dir))
|
|
|
|
model_path = f"{comfy_dir}/models/diffusion_models/{model}"
|
|
vae_path = f"{comfy_dir}/models/diffusion_models/{vae}"
|
|
controlnet_path = f"{comfy_dir}/models/diffusion_models/{controlnet_model}"
|
|
#-----------------------------------------------------------------------
|
|
|
|
# Set up Controlnet-Union-Promax-SDXL model
|
|
config_file = f"{controlnet_path}/config_promax.json"
|
|
config = ControlNetModel_Union.load_config(config_file)
|
|
controlnet_model = ControlNetModel_Union.from_config(config)
|
|
|
|
model_file = f"{controlnet_path}/diffusion_pytorch_model_promax.safetensors"
|
|
state_dict = load_state_dict(model_file)
|
|
|
|
model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
|
|
controlnet_model, state_dict, model_file, f"{controlnet_path}"
|
|
)
|
|
model.to(device="cuda", dtype=torch.float16)
|
|
#-----------------------------------------------------------------------
|
|
|
|
# Set up VAE
|
|
vae = AutoencoderKL.from_pretrained(f"{vae_path}", torch_dtype=torch.float16).to("cuda")
|
|
|
|
if enable_vae_slicing:
|
|
vae.enable_slicing()
|
|
else:
|
|
vae.disable_slicing()
|
|
|
|
if enable_vae_tiling:
|
|
vae.enable_tiling()
|
|
else:
|
|
vae.disable_tiling()
|
|
#-----------------------------------------------------------------------
|
|
|
|
# Load Controlnet + Vae into RealVisXL model
|
|
pipe = StableDiffusionXLFillPipeline.from_pretrained(
|
|
f"{model_path}",
|
|
torch_dtype=torch.float16,
|
|
vae=vae,
|
|
controlnet=controlnet_model,
|
|
variant="fp16"
|
|
)
|
|
|
|
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
|
|
pipe.enable_model_cpu_offload()
|
|
|
|
del model, state_dict, model_file
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
|
|
diffusers_outpaint_pipe = {
|
|
"pipe": pipe,
|
|
"vae": vae,
|
|
"controlnet_model": controlnet_model
|
|
}
|
|
|
|
return (diffusers_outpaint_pipe,)
|
|
|
|
|
|
# Tensor to PIL (grabbed from WAS Suite)
|
|
def tensor2pil(image: torch.Tensor) -> Image.Image:
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
# Convert PIL to Tensor (grabbed from WAS Suite)
|
|
def pil2tensor(image: Image.Image) -> torch.Tensor:
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
class EncodeDiffusersOutpaintPrompt:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"diffusers_outpaint_pipe": ("PIPE", {"tooltip": "Load the diffusers outpaint models."}),
|
|
"extra_prompt": ("STRING", {"default": "", "tooltip": "The extra prompt to append, describing attributes etc. you want to include in the image. Default: \"(extra_prompt), high quality, 4k\""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("PIPE","CONDITIONING",)
|
|
RETURN_NAMES = ("diffusers_outpaint_pipe","diffusers_outpaint_conditioning",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "DiffusersOutpaint"
|
|
|
|
def sample(self, diffusers_outpaint_pipe, extra_prompt=None):
|
|
pipe = diffusers_outpaint_pipe["pipe"]
|
|
|
|
final_prompt = f"{extra_prompt}, high quality, 4k"
|
|
|
|
(prompt_embeds,
|
|
negative_prompt_embeds,
|
|
pooled_prompt_embeds,
|
|
negative_pooled_prompt_embeds,
|
|
) = pipe.encode_prompt(final_prompt)
|
|
|
|
diffusers_outpaint_conditioning = {
|
|
"prompt_embeds": prompt_embeds,
|
|
"negative_prompt_embeds": negative_prompt_embeds,
|
|
"pooled_prompt_embeds": pooled_prompt_embeds,
|
|
"negative_pooled_prompt_embeds": negative_pooled_prompt_embeds
|
|
}
|
|
|
|
del pipe.text_encoder, pipe.text_encoder_2, pipe.tokenizer, pipe.tokenizer_2
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
# Update pipe
|
|
diffusers_outpaint_pipe["pipe"] = pipe
|
|
|
|
return (diffusers_outpaint_pipe,diffusers_outpaint_conditioning,)
|
|
|
|
|
|
class DiffusersImageOutpaint:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"diffusers_outpaint_pipe": ("PIPE", {"tooltip": "Load the diffusers outpaint models."}),
|
|
"diffusers_outpaint_conditioning": ("CONDITIONING", {"tooltip": "The prompt describing what you want."}),
|
|
"diffuser_outpaint_cnet_image": ("IMAGE", {"tooltip": "The image to outpaint."}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Fake seed, workaround used to keep generating different outpaints. Set to -1 to generate different images, or a fixed number to stop that."}),
|
|
"steps": ("INT", {"default": 8, "min": 4, "max": 20, "tooltip": "The number of steps used in the denoising process."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "DiffusersOutpaint"
|
|
|
|
def sample(self, diffusers_outpaint_pipe, diffusers_outpaint_conditioning, diffuser_outpaint_cnet_image, seed, steps):
|
|
# I have to save them here to cache them. The node doesn't load them again
|
|
pipe = diffusers_outpaint_pipe["pipe"]
|
|
vae = diffusers_outpaint_pipe["vae"]
|
|
controlnet_model = diffusers_outpaint_pipe["controlnet_model"]
|
|
prompt_embeds = diffusers_outpaint_conditioning["prompt_embeds"]
|
|
negative_prompt_embeds = diffusers_outpaint_conditioning["negative_prompt_embeds"]
|
|
pooled_prompt_embeds = diffusers_outpaint_conditioning["pooled_prompt_embeds"]
|
|
negative_pooled_prompt_embeds = diffusers_outpaint_conditioning["negative_pooled_prompt_embeds"]
|
|
|
|
cnet_image = diffuser_outpaint_cnet_image
|
|
cnet_image=tensor2pil(cnet_image)
|
|
cnet_image=cnet_image.convert('RGB')
|
|
|
|
generated_images = list(pipe(
|
|
prompt_embeds=prompt_embeds,
|
|
negative_prompt_embeds=negative_prompt_embeds,
|
|
pooled_prompt_embeds=pooled_prompt_embeds,
|
|
negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
|
|
image=cnet_image,
|
|
num_inference_steps=steps
|
|
))
|
|
|
|
del pipe, vae, controlnet_model, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
|
|
last_image = generated_images[-1] # Access the last image
|
|
image = last_image.convert("RGB")
|
|
output=pil2tensor(image)
|
|
|
|
return (output,)
|