237 lines
8.9 KiB
Python
237 lines
8.9 KiB
Python
|
|
import argparse
|
|
import os
|
|
from datetime import datetime
|
|
from typing import Any, Optional
|
|
|
|
from PIL import Image, ImageOps
|
|
|
|
import torch
|
|
|
|
from .ic_custom.pipelines.ic_custom_pipeline import ICCustomPipeline
|
|
import comfy.utils
|
|
import numpy as np
|
|
import folder_paths
|
|
|
|
def ensure_divisible_by_value(image_pil: Image.Image, value: int = 8, interpolate: Image.Resampling = Image.Resampling.LANCZOS) -> Image.Image:
|
|
"""
|
|
Ensure the image dimensions are divisible by value.
|
|
|
|
Args:
|
|
image_pil (Image.Image): The image to ensure divisible by value.
|
|
value (int): The value to ensure divisible by.
|
|
"""
|
|
|
|
w, h = image_pil.size
|
|
w = (w // value) * value
|
|
h = (h // value) * value
|
|
image_pil = image_pil.resize((w, h), interpolate)
|
|
return image_pil
|
|
|
|
def resize_paired_image(
|
|
reference_image: Image.Image,
|
|
target_image: Image.Image,
|
|
mask_target: Image.Image,
|
|
) -> tuple[Image.Image, Image.Image, Image.Image]:
|
|
|
|
ref_w, ref_h = reference_image.size
|
|
target_w, target_h = target_image.size
|
|
|
|
# resize the ref image to the same height as the target image and ensure the ratio remains the same
|
|
if ref_h != target_h:
|
|
scale_ratio = target_h / ref_h
|
|
reference_image = reference_image.resize((int(ref_w * scale_ratio), target_h), interpolate=Image.Resampling.LANCZOS)
|
|
|
|
# Ensure the image dimensions are divisible by 16.
|
|
reference_image = ensure_divisible_by_value(reference_image, value=16, interpolate=Image.Resampling.LANCZOS)
|
|
target_image = ensure_divisible_by_value(target_image, value=16, interpolate=Image.Resampling.LANCZOS)
|
|
mask_target = ensure_divisible_by_value(mask_target, value=16, interpolate=Image.Resampling.NEAREST)
|
|
|
|
return reference_image, target_image, mask_target
|
|
|
|
class RH_ICCustom_Loader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("RHICCustomPipeline",)
|
|
RETURN_NAMES = ("ICCustom Pipeline",)
|
|
FUNCTION = "load"
|
|
|
|
CATEGORY = "Runninghub/ICCustom"
|
|
|
|
def load(self, **kwargs):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
weight_dtype = torch.bfloat16
|
|
|
|
clip_path = os.path.join(folder_paths.models_dir, 'clip_vision', 'clip-vit-large-patch14')
|
|
# t5_path = os.path.join(folder_paths.models_dir, 'clip', 't5-v1_1-xxl')
|
|
t5_path = os.path.join(folder_paths.models_dir, 'clip', 'xflux_text_encoders') #kiki: use xflux instead
|
|
siglip_path = os.path.join(folder_paths.models_dir, 'clip', 'siglip-so400m-patch14-384')
|
|
ae_path = os.path.join(folder_paths.models_dir, 'black-forest-labs', 'FLUX.1-Fill-dev', 'ae.safetensors')
|
|
dit_path = os.path.join(folder_paths.models_dir, 'black-forest-labs', 'FLUX.1-Fill-dev', 'flux1-fill-dev.safetensors')
|
|
redux_path = os.path.join(folder_paths.models_dir, 'IC-Custom', 'flux1-redux-dev.safetensors')
|
|
lora_path = os.path.join(folder_paths.models_dir, 'IC-Custom', 'dit_lora_0x1561.safetensors')
|
|
img_txt_in_path = os.path.join(folder_paths.models_dir, 'IC-Custom', 'dit_txt_img_in_0x1561.safetensors')
|
|
boundary_embeddings_path = os.path.join(folder_paths.models_dir, 'IC-Custom', 'dit_boundary_embeddings_0x1561.safetensors')
|
|
task_register_embeddings_path = os.path.join(folder_paths.models_dir, 'IC-Custom', 'dit_task_register_embeddings_0x1561.safetensors')
|
|
|
|
pipeline = ICCustomPipeline(
|
|
clip_path=clip_path,
|
|
t5_path=t5_path,
|
|
siglip_path=siglip_path,
|
|
ae_path=ae_path,
|
|
dit_path=dit_path,
|
|
redux_path=redux_path,
|
|
lora_path=lora_path,
|
|
img_txt_in_path=img_txt_in_path,
|
|
boundary_embeddings_path=boundary_embeddings_path,
|
|
task_register_embeddings_path=task_register_embeddings_path,
|
|
network_alpha=64,
|
|
double_blocks_idx="0,1,2,3,4,5,6,7,8,9",
|
|
single_blocks_idx="0,1,2,3,4,5,6,7,8,9",
|
|
device=device,
|
|
weight_dtype=weight_dtype,
|
|
offload=True,
|
|
)
|
|
pipeline.set_pipeline_offload(True)
|
|
pipeline.set_show_progress(True)
|
|
|
|
from optimum.quanto import freeze, qint8, quantize, quantization_map, QuantizedDiffusersModel, requantize
|
|
#quantize(pipeline.t5, qint8)
|
|
#freeze(pipeline.t5)
|
|
quantize(pipeline.model, qint8)
|
|
freeze(pipeline.model)
|
|
return (pipeline, )
|
|
# return (None, )
|
|
|
|
class RH_ICCustom_Sampler:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("RHICCustomPipeline", ),
|
|
"ref_image": ("IMAGE", ),
|
|
"prompt": ("STRING", {"multiline": True,
|
|
'default': ''}),
|
|
# "width": ("INT", {"default": 1024}),
|
|
# "height": ("INT", {"default": 1024}),
|
|
"num_inference_steps": ("INT", {"default": 25}),
|
|
"guidance": ("FLOAT", {"default": 40.0}),
|
|
"true_gs": ("FLOAT", {"default": 3.0}),
|
|
"seed": ("INT", {"default": 20, "min": 0, "max": 0xffffffffffffffff,
|
|
"tooltip": "The random seed used for creating the noise."}),
|
|
},
|
|
"optional": {
|
|
"target_image": ("IMAGE", ),
|
|
"target_mask": ("MASK", ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "sample"
|
|
|
|
CATEGORY = "Runninghub/ICCustom"
|
|
|
|
def tensor_2_pil(self, img_tensor):
|
|
if img_tensor is not None:
|
|
i = 255. * img_tensor.squeeze().cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
return img
|
|
else:
|
|
return None
|
|
|
|
def prepare_input_images(self, ref_image, target_image, target_mask):
|
|
img_ref = ref_image
|
|
|
|
# Initialize img_target as a pure white image and mask_target as a pure black image, same size as reference
|
|
img_target = Image.new("RGB", img_ref.size, (255, 255, 255))
|
|
mask_target = Image.new("RGB", img_ref.size, (0, 0, 0))
|
|
|
|
if target_image is not None:
|
|
img_target = target_image
|
|
if target_mask is not None:
|
|
mask_target = target_mask
|
|
|
|
img_ref, img_target, mask_target = resize_paired_image(img_ref, img_target, mask_target)
|
|
|
|
return img_ref, img_target, mask_target
|
|
|
|
def sample(self, **kwargs):
|
|
pipeline = kwargs.get('pipeline')
|
|
# width = kwargs.get('width')
|
|
# height = kwargs.get('height')
|
|
ref_image = self.tensor_2_pil(kwargs.get('ref_image'))
|
|
target_image = self.tensor_2_pil(kwargs.get('target_image'))
|
|
target_mask = self.tensor_2_pil(kwargs.get('target_mask'))
|
|
|
|
if target_mask is None: #pos-free
|
|
target_image = None
|
|
mask_type_ids = 0
|
|
else: #pos-precise
|
|
mask_type_ids = 1
|
|
target_mask = ImageOps.invert(target_mask).convert('RGB')
|
|
|
|
img_ref, img_target, mask_target = self.prepare_input_images(
|
|
ref_image, target_image, target_mask
|
|
)
|
|
img_ip = img_ref.copy()
|
|
cond_w_regions = [img_ref.size[0]]
|
|
|
|
width, height = img_target.size[0] + img_ref.size[0], img_target.size[1]
|
|
prompt = kwargs.get('prompt')
|
|
guidance = kwargs.get('guidance')
|
|
num_steps = kwargs.get('num_inference_steps')
|
|
seed = kwargs.get('seed') % (2 ** 32)
|
|
true_gs = kwargs.get('true_gs')
|
|
|
|
#kiki:hardcode hyperparameters
|
|
neg_prompt = 'worst quality, normal quality, low quality, low res, blurry,'
|
|
self.pbar = comfy.utils.ProgressBar(num_steps)
|
|
|
|
with torch.no_grad():
|
|
image_gen = pipeline(
|
|
prompt=prompt,
|
|
width=width,
|
|
height=height,
|
|
guidance=guidance,
|
|
num_steps=num_steps,
|
|
seed=seed,
|
|
img_ref=img_ref,
|
|
img_target=img_target,
|
|
mask_target=mask_target,
|
|
img_ip=img_ip,
|
|
cond_w_regions=cond_w_regions,
|
|
mask_type_ids=mask_type_ids,
|
|
use_background_preservation=False,
|
|
use_progressive_background_preservation=False,
|
|
background_blend_threshold=0.0,
|
|
true_gs=true_gs,
|
|
neg_prompt=neg_prompt,
|
|
update_func=self.update,
|
|
)[0]
|
|
|
|
image = np.array(image_gen).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
return (image, )
|
|
|
|
def update(self):
|
|
self.pbar.update(1)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RunningHub ICCustom Loader": RH_ICCustom_Loader,
|
|
"RunningHub ICCustom Sampler":RH_ICCustom_Sampler,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RunningHub ICCustom Loader": "RunningHub ICCustom Loader",
|
|
"RunningHub ICCustom Sampler": "RunningHub ICCustom Sampler",
|
|
}
|