352 lines
17 KiB
Python
352 lines
17 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
|
|
import torch
|
|
import os
|
|
from PIL import Image
|
|
import numpy as np
|
|
import sys
|
|
from diffusers import (StableDiffusionXLPipeline, DiffusionPipeline, DDIMScheduler, ControlNetModel,
|
|
KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler,
|
|
AutoPipelineForInpainting, DPMSolverMultistepScheduler, DPMSolverSinglestepScheduler,
|
|
EulerDiscreteScheduler, HeunDiscreteScheduler,
|
|
AutoPipelineForText2Image, StableDiffusionXLControlNetImg2ImgPipeline, KDPM2DiscreteScheduler,
|
|
EulerAncestralDiscreteScheduler, UniPCMultistepScheduler,
|
|
StableDiffusionXLControlNetPipeline)
|
|
from .hidiffusion.hidiffusion import apply_hidiffusion
|
|
import cv2
|
|
import folder_paths
|
|
|
|
dir_path = os.path.dirname(os.path.abspath(__file__))
|
|
path_dir = os.path.dirname(dir_path)
|
|
file_path = os.path.dirname(path_dir)
|
|
|
|
paths=[]
|
|
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):
|
|
if "model_index.json" in files:
|
|
paths.append(os.path.relpath(root, start=search_path))
|
|
# print(paths)
|
|
|
|
scheduler_list = [
|
|
"Euler",
|
|
"Euler a",
|
|
"DDIM",
|
|
"DPM++ 2M",
|
|
"DPM++ 2M Karras",
|
|
"DPM++ 2M SDE",
|
|
"DPM++ 2M SDE Karras",
|
|
"DPM++ SDE",
|
|
"DPM++ SDE Karras",
|
|
"DPM2",
|
|
"DPM2 Karras",
|
|
"DPM2 a",
|
|
"DPM2 a Karras",
|
|
"DPM++ 3M SDE",
|
|
"DPM++ 3M SDE",
|
|
"Heun",
|
|
"LMS",
|
|
"LMS Karras",
|
|
"UniPC",
|
|
"UniPC_Bh2"
|
|
]
|
|
|
|
|
|
def get_sheduler(name):
|
|
scheduler = False
|
|
if name == "Euler":
|
|
scheduler = EulerDiscreteScheduler()
|
|
elif name == "Euler a":
|
|
scheduler = EulerAncestralDiscreteScheduler()
|
|
elif name == "DDIM":
|
|
scheduler = DDIMScheduler()
|
|
elif name == "DPM++ 2M":
|
|
scheduler = DPMSolverMultistepScheduler()
|
|
elif name == "DPM++ 2M Karras":
|
|
scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
|
|
elif name == "DPM++ 2M SDE":
|
|
scheduler = DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++")
|
|
elif name == "DPM++ 2M SDE Karras":
|
|
scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
|
|
elif name == "DPM++ SDE":
|
|
scheduler = DPMSolverSinglestepScheduler()
|
|
elif name == "DPM++ SDE Karras":
|
|
scheduler = DPMSolverSinglestepScheduler(use_karras_sigmas=True)
|
|
elif name == "DPM2":
|
|
scheduler = KDPM2DiscreteScheduler()
|
|
elif name == "DPM2 Karras":
|
|
scheduler = KDPM2DiscreteScheduler(use_karras_sigmas=True)
|
|
elif name == "DPM2 a":
|
|
scheduler = KDPM2AncestralDiscreteScheduler()
|
|
elif name == "DPM2 a Karras":
|
|
scheduler = KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True)
|
|
elif name == "Heun":
|
|
scheduler = HeunDiscreteScheduler()
|
|
elif name == "LMS":
|
|
scheduler = LMSDiscreteScheduler()
|
|
elif name == "LMS Karras":
|
|
scheduler = LMSDiscreteScheduler(use_karras_sigmas=True)
|
|
elif name == "UniPC":
|
|
scheduler = UniPCMultistepScheduler()
|
|
elif name == "UniPC_Bh2":
|
|
scheduler = UniPCMultistepScheduler(solver_type="bh2")
|
|
return scheduler
|
|
|
|
def get_local_path(file_path,model_path):
|
|
path = os.path.join(file_path, "models", "diffusers",model_path)
|
|
model_path = os.path.normpath(path)
|
|
if sys.platform.startswith('win32'):
|
|
model_path = model_path.replace('\\', "/")
|
|
return model_path
|
|
|
|
|
|
class Hidiffusion_Text2Image:
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True,
|
|
"default": "Standing tall amidst the ruins, a stone golem awakens, vines and "
|
|
"flowers sprouting from the crevices in its body"}),
|
|
"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"}),
|
|
"scheduler": (scheduler_list,),
|
|
"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}),
|
|
"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
|
"height": ("INT", {"default": 2048, "min": 64, "max": 8192, "step": 64, "display": "number"}),
|
|
"width": ("INT", {"default": 2048, "min": 64, "max": 8192, "step": 64, "display": "number"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("output_image",)
|
|
FUNCTION = "text2image"
|
|
CATEGORY = "Hidiffusion_Pro"
|
|
|
|
def text2image(self, prompt, negative_prompt,model_local_path, repo_id, scheduler, seed, steps, cfg, eta, height, width):
|
|
|
|
model_path = get_local_path(file_path, model_local_path)
|
|
if repo_id == "":
|
|
repo_id = model_path
|
|
model_type = repo_id.split("/")[-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":
|
|
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
|
pipe = StableDiffusionXLPipeline.from_pretrained(repo_id, scheduler=scheduler,
|
|
torch_dtype=torch.float16, variant="fp16").to("cuda")
|
|
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:
|
|
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")
|
|
else:
|
|
raise "Unsupported model_path or repo_id"
|
|
|
|
# Optional. enable_xformers_memory_efficient_attention can save memory usage and increase inference
|
|
# speed. enable_model_cpu_offload and enable_vae_tiling can save memory usage.
|
|
pipe.enable_xformers_memory_efficient_attention()
|
|
pipe.enable_model_cpu_offload()
|
|
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]
|
|
|
|
output_image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
return (output_image,)
|
|
|
|
|
|
class Hidiffusion_Controlnet_Image:
|
|
|
|
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 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
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"multiline": True,
|
|
"default": "Lara Croft with brown hair, and is wearing a tank top, a brown "
|
|
"backpack."
|
|
"The room is dark and has an old-fashioned decor with a patterned "
|
|
"floor and a wall"
|
|
"featuring a design with arches and a dark area on the right side, "
|
|
"muted color,"
|
|
" high detail, 8k high definition award winning"}),
|
|
"negative_prompt": ("STRING", {"multiline": True,
|
|
"default": "underexposed, poorly drawn hands, duplicate hands, "
|
|
"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,),
|
|
"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_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"}),
|
|
},
|
|
"optional": {
|
|
"mask_image": ("IMAGE",)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("output_image",)
|
|
FUNCTION = "controlnet_image"
|
|
CATEGORY = "Hidiffusion_Pro"
|
|
|
|
|
|
def controlnet_image(self, image, prompt, negative_prompt, model_local_path,repo_id, scheduler, 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):
|
|
|
|
scheduler_used = get_sheduler(scheduler)
|
|
|
|
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)
|
|
scheduler = scheduler_used.from_pretrained(repo_id, subfolder="scheduler")
|
|
model_type = controlnet_repo_id.split("/")[-1]
|
|
|
|
if model_type == "stable-diffusion-xl-1.0-inpainting-0.1":
|
|
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 == "controlnet-canny-sdxl-1.0":
|
|
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)
|
|
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]
|
|
|
|
elif controlnet_type == "text2img":
|
|
canny_image = self.get_canny_image(ori_image, canny_minval, canny_maxval)
|
|
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]
|
|
else:
|
|
raise "Unsupported model"
|
|
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,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Hidiffusion_Text2Image": Hidiffusion_Text2Image,
|
|
"Hidiffusion_Controlnet_Image": Hidiffusion_Controlnet_Image
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Hidiffusion_Text2Image": "Hidiffusion_Text2Image",
|
|
"Hidiffusion_Controlnet_Image": "Hidiffusion_Controlnet_Image"
|
|
}
|