Update Hidiffusion_node.py
update
This commit is contained in:
smthemex
+280
-232
@@ -11,20 +11,21 @@ from diffusers import (StableDiffusionXLPipeline, DiffusionPipeline, DDIMSchedul
|
||||
AutoPipelineForInpainting, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler,
|
||||
EulerDiscreteScheduler, HeunDiscreteScheduler, UNet2DConditionModel,
|
||||
AutoPipelineForText2Image, StableDiffusionXLControlNetImg2ImgPipeline, KDPM2DiscreteScheduler,
|
||||
EulerAncestralDiscreteScheduler, UniPCMultistepScheduler,
|
||||
EulerAncestralDiscreteScheduler, UniPCMultistepScheduler, AutoencoderKL,
|
||||
StableDiffusionXLControlNetPipeline, DDPMScheduler, TCDScheduler, LCMScheduler)
|
||||
|
||||
from .hidiffusion.hidiffusion import apply_hidiffusion
|
||||
import cv2
|
||||
import folder_paths
|
||||
from safetensors.torch import load_file
|
||||
|
||||
|
||||
dir_path = os.path.dirname(os.path.abspath(__file__))
|
||||
path_dir = os.path.dirname(dir_path)
|
||||
file_path = os.path.dirname(path_dir)
|
||||
|
||||
paths = []
|
||||
paths_a = []
|
||||
paths_unet = []
|
||||
for search_path in folder_paths.get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
for root, subdir, files in os.walk(search_path, followlinks=True):
|
||||
@@ -32,10 +33,10 @@ for search_path in folder_paths.get_folder_paths("diffusers"):
|
||||
paths.append(os.path.relpath(root, start=search_path))
|
||||
if "config.json" in files:
|
||||
paths_a.append(os.path.relpath(root, start=search_path))
|
||||
paths_a = [z for z in paths_a if "controlnet-canny-sdxl-1.0" in z] + [p for p in paths_a if
|
||||
"MistoLine" in p] + [o for o in
|
||||
paths_a if
|
||||
"lcm-sdxl" in o]
|
||||
paths_a = ([
|
||||
z for z in paths_a if "controlnet-canny-sdxl-1.0" in z]
|
||||
+ [p for p in paths_a if "MistoLine" in p]
|
||||
+ [o for o in paths_a if "lcm-sdxl" in o] + [Q for Q in paths_a if "controlnet-openpose-sdxl-1.0" in Q])
|
||||
|
||||
if paths != [] or paths_a != []:
|
||||
paths = [] + [x for x in paths if x] + [y for y in paths_a if y]
|
||||
@@ -62,7 +63,6 @@ scheduler_list = [
|
||||
"LMS",
|
||||
"LMS Karras",
|
||||
"UniPC",
|
||||
"UniPC_Bh2",
|
||||
"TCD"
|
||||
]
|
||||
|
||||
@@ -72,8 +72,12 @@ sdxl_lightning_list = [
|
||||
"sdxl_lightning_4step_unet.safetensors",
|
||||
"sdxl_lightning_8step_unet.safetensors"
|
||||
"Hyper-SDXL-1step-Unet.safetensors",
|
||||
"lcm-sdxl-base-1.0.safetensors",
|
||||
]
|
||||
|
||||
normal_model_list = ["stable-diffusion-2-1-base", "stable-diffusion-v1-5", "playground-v2.5-1024px-aesthetic",
|
||||
"Ghibli-Diffusion", "playground-v2-1024px-aesthetic"]
|
||||
|
||||
|
||||
def get_sheduler(name):
|
||||
scheduler = False
|
||||
@@ -107,21 +111,26 @@ def get_sheduler(name):
|
||||
scheduler = KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True)
|
||||
elif name == "Heun":
|
||||
scheduler = HeunDiscreteScheduler()
|
||||
elif name == "LCD":
|
||||
elif name == "LCM":
|
||||
scheduler = LCMScheduler()
|
||||
elif name == "LMS":
|
||||
scheduler = LMSDiscreteScheduler()
|
||||
elif name == "LMS Karras":
|
||||
scheduler = LMSDiscreteScheduler(use_karras_sigmas=True)
|
||||
elif name == "UniPC_Bh1":
|
||||
scheduler = UniPCMultistepScheduler(solver_type="bh1")
|
||||
elif name == "UniPC_Bh2":
|
||||
scheduler = UniPCMultistepScheduler(solver_type="bh2")
|
||||
elif name == "UniPC":
|
||||
scheduler = UniPCMultistepScheduler()
|
||||
elif name == "TCD":
|
||||
scheduler = TCDScheduler()
|
||||
return scheduler
|
||||
|
||||
|
||||
def tensor_to_image(tensor):
|
||||
tensor = tensor.cpu()
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
|
||||
|
||||
def get_local_path(file_path, model_path):
|
||||
path = os.path.join(file_path, "models", "diffusers", model_path)
|
||||
model_path = os.path.normpath(path)
|
||||
@@ -130,8 +139,46 @@ def get_local_path(file_path, model_path):
|
||||
return model_path
|
||||
|
||||
|
||||
class Hidiffusion_Text2Image:
|
||||
def get_instance_path(path):
|
||||
instance_path = os.path.normpath(path)
|
||||
if sys.platform.startswith('win32'):
|
||||
instance_path = instance_path.replace('\\', "/")
|
||||
return instance_path
|
||||
|
||||
|
||||
class Diffusers_Or_Repo_Choice:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"local_model_path": (paths,),
|
||||
"repo_id": ("STRING", {"default": "stabilityai/stable-diffusion-xl-base-1.0"})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("repo_id",)
|
||||
FUNCTION = "repo_choice"
|
||||
CATEGORY = "Hidiffusion_Pro"
|
||||
|
||||
def repo_choice(self, local_model_path, repo_id):
|
||||
if repo_id == "":
|
||||
if local_model_path == ["no model in default diffusers directory", ]:
|
||||
raise "you need fill repo_id or download model in diffusers directory "
|
||||
elif local_model_path != ["no model in default diffusers directory", ]:
|
||||
model_path = get_local_path(file_path, local_model_path)
|
||||
repo_id = get_instance_path(model_path)
|
||||
elif repo_id != "" and repo_id.find("/") == -1:
|
||||
raise "Incorrect repo_id format"
|
||||
elif repo_id != "" and repo_id.find("\\") != -1:
|
||||
repo_id = get_instance_path(repo_id)
|
||||
return (repo_id,)
|
||||
|
||||
|
||||
class Hi_Text2Img:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@@ -145,10 +192,9 @@ class Hidiffusion_Text2Image:
|
||||
"negative_prompt": ("STRING", {"multiline": True,
|
||||
"default": "blurry, ugly, duplicate, poorly drawn face, deformed, "
|
||||
"mosaic, artifacts, bad limbs"}),
|
||||
"model_local_path": (paths,),
|
||||
"repo_id": ("STRING", {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
|
||||
"repo_id": ("STRING", {"forceInput": True}),
|
||||
"scheduler": (scheduler_list,),
|
||||
"sdxl_lightning_model": ([] + folder_paths.get_filename_list("unet"),),
|
||||
"unet_model": ([] + folder_paths.get_filename_list("unet"),),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 50, "min": 1, "max": 10000}),
|
||||
"cfg": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
||||
@@ -163,47 +209,24 @@ class Hidiffusion_Text2Image:
|
||||
FUNCTION = "text2image"
|
||||
CATEGORY = "Hidiffusion_Pro"
|
||||
|
||||
def text2image(self, prompt, negative_prompt, model_local_path, repo_id, scheduler, sdxl_lightning_model, seed,
|
||||
steps, cfg, eta, height,
|
||||
width):
|
||||
if model_local_path == ["no model in default diffusers directory", ] and repo_id == "":
|
||||
raise "you need fill repo_id or download model in diffusers dir "
|
||||
|
||||
model_path = get_local_path(file_path, model_local_path)
|
||||
if repo_id == "":
|
||||
repo_id = model_path
|
||||
model_type = repo_id.split("/")[-1]
|
||||
model_type_lcm = model_path.split("/")[-1]
|
||||
def text2image(self, prompt, negative_prompt, repo_id, scheduler, unet_model, seed,
|
||||
steps, cfg, eta, height, width):
|
||||
model_type = repo_id.rsplit("/")[-1]
|
||||
scheduler_used = get_sheduler(scheduler)
|
||||
model_list = ["stable-diffusion-2-1-base", "stable-diffusion-v1-5", "Ghibli-Diffusion"]
|
||||
if model_type == "stable-diffusion-xl-base-1.0":
|
||||
if sdxl_lightning_model in sdxl_lightning_list:
|
||||
base = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
# repo = "ByteDance/SDXL-Lightning"
|
||||
# ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
|
||||
light_path = os.path.join(file_path, "models", "unet", sdxl_lightning_model)
|
||||
light_model_path = os.path.normpath(light_path)
|
||||
if sys.platform.startswith('win32'):
|
||||
light_model_path = light_model_path.replace('\\', "/")
|
||||
ckpt = light_model_path
|
||||
# print(ckpt)
|
||||
# Load model.
|
||||
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"))
|
||||
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16,
|
||||
if unet_model in sdxl_lightning_list:
|
||||
light_path = os.path.join(file_path, "models", "unet", unet_model)
|
||||
ckpt = get_instance_path(light_path)
|
||||
# UNet2DConditionModel.load_config(repo_id, subfolder="unet")
|
||||
unet = UNet2DConditionModel.from_pretrained(repo_id, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"), strict=False, )
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(repo_id, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
elif model_type_lcm == "lcm-sdxl":
|
||||
base = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
# print(model_path)
|
||||
unet = UNet2DConditionModel.from_pretrained(model_path, state_dict=False, torch_dtype=torch.float16,
|
||||
variant="fp16")
|
||||
pipe = DiffusionPipeline.from_pretrained(base, unet=unet,
|
||||
torch_dtype=torch.float16, variant="fp16")
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
|
||||
pipe.to("cuda")
|
||||
if unet_model == "lcm-sdxl-base-1.0.safetensors" or unet_model == "Hyper-SDXL-1step-Unet.safetensors":
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(repo_id, scheduler=scheduler,
|
||||
@@ -211,12 +234,7 @@ class Hidiffusion_Text2Image:
|
||||
elif model_type == "sdxl-turbo":
|
||||
pipe = AutoPipelineForText2Image.from_pretrained(repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to('cuda')
|
||||
elif model_type == "playground-v2-1024px-aesthetic":
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = DiffusionPipeline.from_pretrained(repo_id, scheduler=scheduler,
|
||||
torch_dtype=torch.float16, use_safetensors=True,
|
||||
add_watermarker=False, variant="fp16").to("cuda")
|
||||
elif model_type in model_list:
|
||||
elif model_type in normal_model_list:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = DiffusionPipeline.from_pretrained(repo_id, scheduler=scheduler,
|
||||
torch_dtype=torch.float16, variant="fp16").to("cuda")
|
||||
@@ -230,42 +248,178 @@ class Hidiffusion_Text2Image:
|
||||
pipe.enable_vae_tiling()
|
||||
# Apply hidiffusion with a single line of code.
|
||||
apply_hidiffusion(pipe)
|
||||
if steps > 0 and model_type == "sdxl-turbo":
|
||||
image = pipe(prompt, num_inference_steps=steps, guidance_scale=cfg, height=height,
|
||||
width=width, eta=eta, seed=seed,
|
||||
negative_prompt=negative_prompt).images[0]
|
||||
else:
|
||||
image = pipe(prompt, guidance_scale=cfg, height=height, width=width, eta=eta, num_inference_steps=steps,
|
||||
seed=seed,
|
||||
negative_prompt=negative_prompt).images[0]
|
||||
|
||||
image = pipe(prompt, num_inference_steps=steps, guidance_scale=cfg, height=height, width=width, eta=eta,
|
||||
seed=seed, negative_prompt=negative_prompt).images[0]
|
||||
output_image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return (output_image,)
|
||||
|
||||
|
||||
class Hidiffusion_Controlnet_Image:
|
||||
class Hi_Control2Img:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def tensor_to_image(self, tensor):
|
||||
tensor = tensor.cpu()
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
def inpainting_gener(self, image, repo_id, controlnet_repo_id, prompt, negative_prompt, scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta, controlnet_strength, height, width, mask_image):
|
||||
scheduler_used = get_sheduler(scheduler)
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
if unet_model in sdxl_lightning_list:
|
||||
light_path = os.path.join(file_path, "models", "unet", unet_model)
|
||||
ckpt = get_instance_path(light_path)
|
||||
# UNet2DConditionModel.load_config(repo_id, subfolder="unet")
|
||||
unet = UNet2DConditionModel.from_pretrained(repo_id, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"), strict=False, )
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(repo_id, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
if unet_model == "lcm-sdxl-base-1.0.safetensors" or unet_model == "Hyper-SDXL-1step-Unet.safetensors":
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
pipe = AutoPipelineForInpainting.from_pretrained(
|
||||
controlnet_repo_id, torch_dtype=torch.float16, variant="fp16",
|
||||
scheduler=scheduler
|
||||
)
|
||||
apply_hidiffusion(pipe)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.enable_vae_tiling()
|
||||
|
||||
def get_canny_image(self, img, a, b): # get canny image
|
||||
img = np.array(img)
|
||||
img = cv2.Canny(img, a, b)
|
||||
img = img[:, :, None]
|
||||
img = np.concatenate([img, img, img], axis=2)
|
||||
canny_image = Image.fromarray(img)
|
||||
return canny_image
|
||||
output_img = pipe(prompt=prompt, image=image, mask_image=mask_image, height=height, width=width,
|
||||
strength=controlnet_strength, num_inference_steps=steps, seed=seed,
|
||||
guidance_scale=cfg, negative_prompt=negative_prompt, eta=eta).images[0]
|
||||
return output_img
|
||||
|
||||
def control_img2img(self, image, prompt, negative_prompt, repo_id, controlnet_repo_id, scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta, controlnet_conditioning_scale,
|
||||
controlnet_strength, height, width, control_image):
|
||||
scheduler_used = get_sheduler(scheduler)
|
||||
controlnet = ControlNetModel.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
if unet_model in sdxl_lightning_list:
|
||||
light_path = os.path.join(file_path, "models", "unet", unet_model)
|
||||
ckpt = get_instance_path(light_path)
|
||||
|
||||
# UNet2DConditionModel.load_config(repo_id, subfolder="unet")
|
||||
unet = UNet2DConditionModel.from_pretrained(repo_id, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"), strict=False, )
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(repo_id, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
if unet_model == "lcm-sdxl-base-1.0.safetensors" or unet_model == "Hyper-SDXL-1step-Unet.safetensors":
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(repo_id, controlnet=controlnet,
|
||||
scheduler=scheduler,
|
||||
torch_dtype=torch.float16, ).to("cuda")
|
||||
apply_hidiffusion(pipe)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.enable_vae_tiling()
|
||||
|
||||
output_img = pipe(prompt=prompt,
|
||||
image=image,
|
||||
control_image=control_image,
|
||||
height=height,
|
||||
width=width,
|
||||
seed=seed,
|
||||
strength=controlnet_strength,
|
||||
num_inference_steps=steps,
|
||||
controlnet_conditioning_scale=controlnet_conditioning_scale,
|
||||
guidance_scale=cfg,
|
||||
negative_prompt=negative_prompt,
|
||||
eta=eta
|
||||
).images[0]
|
||||
return output_img
|
||||
|
||||
def control_txt2img(self, prompt, negative_prompt, repo_id, controlnet_repo_id, scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta, controlnet_conditioning_scale,
|
||||
controlnet_strength, height, width, control_image):
|
||||
scheduler_used = get_sheduler(scheduler)
|
||||
controlnet = ControlNetModel.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
if unet_model in sdxl_lightning_list:
|
||||
light_path = os.path.join(file_path, "models", "unet", unet_model)
|
||||
ckpt = get_instance_path(light_path)
|
||||
# UNet2DConditionModel.load_config(repo_id, subfolder="unet")
|
||||
unet = UNet2DConditionModel.from_config(repo_id, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"), strict=False, )
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(repo_id, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
if unet_model == "lcm-sdxl-base-1.0.safetensors" or unet_model == "Hyper-SDXL-1step-Unet.safetensors":
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
|
||||
repo_id, controlnet=controlnet, torch_dtype=torch.float16,
|
||||
scheduler=scheduler
|
||||
)
|
||||
apply_hidiffusion(pipe)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.enable_vae_tiling()
|
||||
|
||||
output_img = pipe(
|
||||
prompt=prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=control_image,
|
||||
height=height, width=width, guidance_scale=cfg, negative_prompt=negative_prompt, seed=seed,
|
||||
strength=controlnet_strength,
|
||||
num_inference_steps=steps, eta=eta
|
||||
).images[0]
|
||||
return output_img
|
||||
|
||||
def openpose_gener(self, prompt, negative_prompt, repo_id, controlnet_repo_id, scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta, controlnet_conditioning_scale,
|
||||
controlnet_strength, height, width, mask_image):
|
||||
scheduler_used = get_sheduler(scheduler)
|
||||
controlnet = ControlNetModel.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
|
||||
if unet_model in sdxl_lightning_list:
|
||||
light_path = os.path.join(file_path, "models", "unet", unet_model)
|
||||
ckpt = get_instance_path(light_path)
|
||||
# UNet2DConditionModel.load_config(repo_id, subfolder="unet")
|
||||
unet = UNet2DConditionModel.from_config(repo_id, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"), strict=False, )
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(repo_id, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
if unet_model == "lcm-sdxl-base-1.0.safetensors" or unet_model == "Hyper-SDXL-1step-Unet.safetensors":
|
||||
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
|
||||
repo_id, controlnet=controlnet, vae=vae, torch_dtype=torch.float16,
|
||||
scheduler=scheduler
|
||||
)
|
||||
|
||||
apply_hidiffusion(pipe)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.enable_vae_tiling()
|
||||
|
||||
output_img = pipe(
|
||||
prompt=prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=mask_image,
|
||||
height=height, width=width, guidance_scale=cfg, negative_prompt=negative_prompt, seed=seed,
|
||||
strength=controlnet_strength,
|
||||
num_inference_steps=steps, eta=eta
|
||||
).images[0]
|
||||
return output_img
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"repo_id": ("STRING", {"forceInput": True}),
|
||||
"controlnet_repo_id": ("STRING", {"forceInput": True}),
|
||||
"prompt": ("STRING", {"multiline": True,
|
||||
"default": "Lara Croft with brown hair, and is wearing a tank top, a brown "
|
||||
"backpack."
|
||||
@@ -279,186 +433,80 @@ class Hidiffusion_Controlnet_Image:
|
||||
"overexposed,"
|
||||
"bad art, beginner, amateur, abstract, disfigured, "
|
||||
"deformed, close up, weird colors, watermark"}),
|
||||
"model_local_path": (paths,),
|
||||
"repo_id": ("STRING", {"default": "stabilityai/stable-diffusion-xl-base-1.0"}),
|
||||
|
||||
"scheduler": (scheduler_list,),
|
||||
"sdxl_lightning_model": ([] + folder_paths.get_filename_list("unet"),),
|
||||
"unet_model": ([] + folder_paths.get_filename_list("unet"),),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||||
"controlnet_local_path": (paths,),
|
||||
"controlnet_repo_id": ("STRING", {"default": "diffusers/controlnet-canny-sdxl-1.0"}),
|
||||
"cfg": ("FLOAT", {"default": 12.5, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
||||
"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
||||
"canny_minval": ("INT", {"default": 100, "min": 0, "max": 255}),
|
||||
"canny_maxval": ("INT", {"default": 200, "min": 0, "max": 255}),
|
||||
"controlnet_type": (["text2img", "img2img", ],),
|
||||
"controlnet_choice": (["text2img", "img2img", ],),
|
||||
"controlnet_conditioning_scale": (
|
||||
"FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.1, "round": 0.01}),
|
||||
"controlnet_strength": (
|
||||
"FLOAT", {"default": 0.99, "min": 0.1, "max": 1.0, "step": 0.1, "round": 0.01}),
|
||||
"height": ("INT", {"default": 1536, "min": 64, "max": 8192, "step": 4, "display": "number"}),
|
||||
"width": ("INT", {"default": 2048, "min": 64, "max": 8192, "step": 4, "display": "number"}),
|
||||
"width": ("INT", {"default": 2048, "min": 64, "max": 8192, "step": 4, "display": "number"})
|
||||
},
|
||||
"optional": {
|
||||
"control_image": ("IMAGE",),
|
||||
"mask_image": ("IMAGE",)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("output_image",)
|
||||
FUNCTION = "controlnet_image"
|
||||
FUNCTION = "xl_controlnet_image"
|
||||
CATEGORY = "Hidiffusion_Pro"
|
||||
|
||||
def controlnet_image(self, image, prompt, negative_prompt, model_local_path, repo_id, scheduler,
|
||||
sdxl_lightning_model,
|
||||
controlnet_local_path, controlnet_repo_id, seed, steps, cfg,
|
||||
eta, canny_minval, canny_maxval, controlnet_type,
|
||||
controlnet_conditioning_scale, controlnet_strength, height, width,
|
||||
mask_image):
|
||||
def xl_controlnet_image(self, image, repo_id, controlnet_repo_id, prompt, negative_prompt, scheduler,
|
||||
unet_model, seed, steps, cfg,
|
||||
eta, controlnet_choice,
|
||||
controlnet_conditioning_scale, controlnet_strength, height, width, control_image,
|
||||
mask_image, ):
|
||||
|
||||
scheduler_used = get_sheduler(scheduler)
|
||||
if model_local_path == ["no model in default diffusers directory", ] and repo_id == "":
|
||||
raise "you need fill repo_id or download model in diffusers dir "
|
||||
model_path = get_local_path(file_path, model_local_path)
|
||||
controlnet_local_path = get_local_path(file_path, controlnet_local_path)
|
||||
if repo_id == "":
|
||||
repo_id = model_path
|
||||
if controlnet_repo_id == "":
|
||||
controlnet_repo_id = controlnet_local_path
|
||||
# print(model_path,controlnet_local_path)
|
||||
ori_image = self.tensor_to_image(image)
|
||||
mask = self.tensor_to_image(mask_image)
|
||||
model_type = controlnet_repo_id.split("/")[-1]
|
||||
image = tensor_to_image(image)
|
||||
control_image = np.transpose(control_image, (
|
||||
0, 3, 1, 2)) # 将comfy输入的（batch_size,channels,rows,cols）改成条件需求的（3通道，16输出，3，1 ，pading=1）
|
||||
mask_image = tensor_to_image(mask_image)
|
||||
model_type = repo_id.rsplit("/")[-1]
|
||||
control_model_type = controlnet_repo_id.rsplit("/")[-1]
|
||||
controlnet_suport = ["controlnet-canny-sdxl-1.0", "MistoLine"]
|
||||
if model_type == "stable-diffusion-xl-1.0-inpainting-0.1":
|
||||
if sdxl_lightning_model in sdxl_lightning_list:
|
||||
base = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
# repo = "ByteDance/SDXL-Lightning"
|
||||
# ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
|
||||
light_path = os.path.join(file_path, "models", "unet", sdxl_lightning_model)
|
||||
light_model_path = os.path.normpath(light_path)
|
||||
if sys.platform.startswith('win32'):
|
||||
light_model_path = light_model_path.replace('\\', "/")
|
||||
ckpt = light_model_path
|
||||
# print(ckpt)
|
||||
# Load model.
|
||||
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"))
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
|
||||
if control_model_type == "stable-diffusion-xl-1.0-inpainting-0.1":
|
||||
output_img = self.inpainting_gener(image, repo_id, controlnet_repo_id, prompt, negative_prompt, scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta, controlnet_strength,
|
||||
height, width, mask_image)
|
||||
elif control_model_type == "controlnet-openpose-sdxl-1.0":
|
||||
output_img = self.openpose_gener(prompt, negative_prompt, repo_id, controlnet_repo_id, scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta, controlnet_conditioning_scale,
|
||||
controlnet_strength, height, width, mask_image)
|
||||
elif control_model_type in controlnet_suport:
|
||||
if controlnet_choice == "img2img":
|
||||
output_img = self.control_img2img(image, prompt, negative_prompt, repo_id, controlnet_repo_id,
|
||||
scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta,
|
||||
controlnet_conditioning_scale,
|
||||
controlnet_strength, height, width, control_image)
|
||||
else:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = AutoPipelineForInpainting.from_pretrained(
|
||||
controlnet_repo_id, torch_dtype=torch.float16, variant="fp16",
|
||||
scheduler=scheduler
|
||||
)
|
||||
apply_hidiffusion(pipe)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.enable_vae_tiling()
|
||||
|
||||
output_img = pipe(prompt=prompt, image=ori_image, mask_image=mask, height=height, width=width,
|
||||
strength=controlnet_strength, num_inference_steps=steps, seed=seed,
|
||||
guidance_scale=cfg, negative_prompt=negative_prompt, eta=eta).images[0]
|
||||
elif model_type in controlnet_suport:
|
||||
controlnet = ControlNetModel.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to(
|
||||
"cuda")
|
||||
if controlnet_type == "img2img":
|
||||
# get canny image
|
||||
canny_image = self.get_canny_image(ori_image, canny_minval, canny_maxval)
|
||||
if sdxl_lightning_model in sdxl_lightning_list:
|
||||
base = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
# repo = "ByteDance/SDXL-Lightning"
|
||||
# ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
|
||||
light_path = os.path.join(file_path, "models", "unet", sdxl_lightning_model)
|
||||
light_model_path = os.path.normpath(light_path)
|
||||
if sys.platform.startswith('win32'):
|
||||
light_model_path = light_model_path.replace('\\', "/")
|
||||
ckpt = light_model_path
|
||||
# print(ckpt)
|
||||
# Load model.
|
||||
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"))
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(repo_id, controlnet=controlnet,
|
||||
scheduler=scheduler,
|
||||
torch_dtype=torch.float16, ).to(
|
||||
"cuda")
|
||||
apply_hidiffusion(pipe)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.enable_vae_tiling()
|
||||
|
||||
output_img = pipe(prompt=prompt,
|
||||
image=ori_image,
|
||||
control_image=canny_image,
|
||||
height=height,
|
||||
width=width,
|
||||
seed=seed,
|
||||
strength=controlnet_strength,
|
||||
num_inference_steps=steps,
|
||||
controlnet_conditioning_scale=controlnet_conditioning_scale,
|
||||
guidance_scale=cfg,
|
||||
negative_prompt=negative_prompt,
|
||||
eta=eta
|
||||
).images[0]
|
||||
else:
|
||||
canny_image = self.get_canny_image(ori_image, canny_minval, canny_maxval)
|
||||
if sdxl_lightning_model in sdxl_lightning_list:
|
||||
base = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
# repo = "ByteDance/SDXL-Lightning"
|
||||
# ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
|
||||
light_path = os.path.join(file_path, "models", "unet", sdxl_lightning_model)
|
||||
light_model_path = os.path.normpath(light_path)
|
||||
if sys.platform.startswith('win32'):
|
||||
light_model_path = light_model_path.replace('\\', "/")
|
||||
ckpt = light_model_path
|
||||
# print(ckpt)
|
||||
# Load model.
|
||||
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(ckpt, device="cuda"))
|
||||
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
# Ensure sampler uses "trailing" timesteps.
|
||||
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config, timestep_spacing="trailing")
|
||||
else:
|
||||
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
||||
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
|
||||
repo_id, controlnet=controlnet, torch_dtype=torch.float16,
|
||||
scheduler=scheduler
|
||||
)
|
||||
apply_hidiffusion(pipe)
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
pipe.enable_model_cpu_offload()
|
||||
pipe.enable_vae_tiling()
|
||||
|
||||
output_img = pipe(
|
||||
prompt=prompt, controlnet_conditioning_scale=controlnet_conditioning_scale, image=canny_image,
|
||||
seed=seed,
|
||||
height=height, width=width, guidance_scale=cfg, negative_prompt=negative_prompt,
|
||||
num_inference_steps=steps, eta=eta
|
||||
).images[0]
|
||||
output_img = self.control_txt2img(prompt, negative_prompt, repo_id, controlnet_repo_id, scheduler,
|
||||
unet_model, model_type, seed, steps, cfg, eta,
|
||||
controlnet_conditioning_scale,
|
||||
controlnet_strength, height, width, control_image)
|
||||
else:
|
||||
raise "Unsupported model_path or repo_id"
|
||||
|
||||
output_image = torch.from_numpy(np.array(output_img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return (output_image,)
|
||||
output_img = torch.from_numpy(np.array(output_img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return (output_img,)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Hidiffusion_Text2Image": Hidiffusion_Text2Image,
|
||||
"Hidiffusion_Controlnet_Image": Hidiffusion_Controlnet_Image
|
||||
"Diffusers_Or_Repo_Choice": Diffusers_Or_Repo_Choice,
|
||||
"Hi_Text2Img": Hi_Text2Img,
|
||||
"Hi_Control2Img": Hi_Control2Img
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Hidiffusion_Text2Image": "Hidiffusion_Text2Image",
|
||||
"Hidiffusion_Controlnet_Image": "Hidiffusion_Controlnet_Image"
|
||||
"Diffusers_Or_Repo_Choice": "Diffusers_Or_Repo_Choice",
|
||||
"Hi_Text2Img": "Hi_Text2Img",
|
||||
"Hi_Control2Img": "Hi_Control2Img"
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user