Update Hidiffusion_node.py
change
This commit is contained in:
smthemex
+66
-81
@@ -17,7 +17,7 @@ from diffusers import (StableDiffusionXLPipeline, DiffusionPipeline, DDIMSchedul
|
||||
StableDiffusionPipeline, StableDiffusionControlNetPipeline, StableDiffusionXLInpaintPipeline)
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
from .hidiffusion.hidiffusion import apply_hidiffusion
|
||||
from .hidiffusion.hidiffusion import apply_hidiffusion,remove_hidiffusion
|
||||
import folder_paths
|
||||
from safetensors.torch import load_file, load
|
||||
import yaml
|
||||
@@ -208,7 +208,8 @@ class HI_Diffusers_Model_Loader:
|
||||
"ip_adapter": ("BOOLEAN", {"default": False},),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
RETURN_TYPES = ("MODEL", "STRING",)
|
||||
RETURN_NAMES = ("model", "model_info",)
|
||||
FUNCTION = "loader_models"
|
||||
@@ -280,33 +281,22 @@ class HI_Diffusers_Model_Loader:
|
||||
if not controlnet_repo_id:
|
||||
original_config_file = os.path.join(dir_path, "weights", "sd_xl_base.yaml")
|
||||
if controlnet_model!="none" and "inpaint" in controlnet_model :
|
||||
unet = UNet2DConditionModel.from_config("stabilityai/stable-diffusion-xl-base-1.0",
|
||||
subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(controlnet_single_model, device="cuda"), strict=False, )
|
||||
if dif_version_int >= 28:
|
||||
model = StableDiffusionXLInpaintPipeline.from_single_file(ckpt_path,unet=unet,
|
||||
model = StableDiffusionXLInpaintPipeline.from_single_file(ckpt_path,
|
||||
original_config=original_config_file,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLInpaintPipeline.from_single_file(ckpt_path,unet=unet,
|
||||
model = StableDiffusionXLInpaintPipeline.from_single_file(ckpt_path,
|
||||
original_config_file=original_config_file,
|
||||
torch_dtype=torch.float16,variant="fp16").to("cuda")
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
else:
|
||||
if dif_version_int >= 28:
|
||||
model = StableDiffusionXLInpaintPipeline.from_single_file(ckpt_path,
|
||||
original_config=original_config_file,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLInpaintPipeline.from_single_file(ckpt_path,
|
||||
original_config_file=original_config_file,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
|
||||
raise "need inpaiting model"
|
||||
else:
|
||||
model = AutoPipelineForInpainting.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
|
||||
if unet_model in sdxl_lightning_list:
|
||||
ckpt = folder_paths.get_full_path("unet", unet_model)
|
||||
unet = UNet2DConditionModel.from_config("stabilityai/stable-diffusion-xl-base-1.0",
|
||||
@@ -320,72 +310,60 @@ class HI_Diffusers_Model_Loader:
|
||||
model_type = "stable-diffusion-xl-base-1.0"
|
||||
if not repo_id:
|
||||
original_config_file = os.path.join(dir_path, "weights", "sd_xl_base.yaml")
|
||||
if control_model_type == "none": #不启用
|
||||
if dif_version_int >= 28:
|
||||
model = StableDiffusionXLPipeline.from_single_file(
|
||||
ckpt_path, original_config=original_config_file, torch_dtype=torch.float16).to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLPipeline.from_single_file(
|
||||
ckpt_path, original_config_file=original_config_file, torch_dtype=torch.float16).to("cuda")
|
||||
|
||||
if dif_version_int >= 28:
|
||||
model = StableDiffusionXLPipeline.from_single_file(
|
||||
ckpt_path, original_config=original_config_file, torch_dtype=torch.float16).to("cuda")
|
||||
else:
|
||||
if controlnet_single_model:
|
||||
if not controlnet_repo_id:
|
||||
unet = UNet2DConditionModel.from_config("stabilityai/stable-diffusion-xl-base-1.0",
|
||||
subfolder="unet").to("cuda", torch.float16)
|
||||
unet.load_state_dict(load_file(controlnet_single_model, device="cuda"), strict=False, )
|
||||
controlnet = ControlNetModel.from_unet(unet).to("cuda")
|
||||
else:
|
||||
controlnet = ControlNetModel.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
model = StableDiffusionXLPipeline.from_single_file(
|
||||
ckpt_path, original_config_file=original_config_file, torch_dtype=torch.float16).to("cuda")
|
||||
if control_model_type!="none":
|
||||
if not controlnet_repo_id:
|
||||
controlnet = ControlNetModel.from_unet(model.unet)
|
||||
cn_state_dict = load_file( controlnet_single_model, device="cuda")
|
||||
controlnet.load_state_dict(cn_state_dict, strict=False)
|
||||
|
||||
else:
|
||||
controlnet = ""
|
||||
if dif_version_int >= 28:
|
||||
if function_choice == "img2img":
|
||||
model = StableDiffusionXLControlNetImg2ImgPipeline.from_single_file(ckpt_path,original_config=original_config_file,
|
||||
controlnet=controlnet,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to(
|
||||
"cuda")
|
||||
else:
|
||||
model = StableDiffusionXLControlNetPipeline.from_single_file(
|
||||
ckpt_path, original_config=original_config_file, controlnet=controlnet,
|
||||
torch_dtype=torch.float16).to("cuda")
|
||||
else:
|
||||
if function_choice == "img2img":
|
||||
model = StableDiffusionXLControlNetImg2ImgPipeline.from_single_file(ckpt_path,
|
||||
original_config_file=original_config_file,
|
||||
controlnet=controlnet,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLControlNetPipeline.from_single_file(
|
||||
ckpt_path, original_config_file=original_config_file,controlnet=controlnet, torch_dtype=torch.float16).to("cuda")
|
||||
else:
|
||||
if control_model_type == "none":
|
||||
if function_choice == "img2img":
|
||||
model = StableDiffusionXLImg2ImgPipeline.from_pretrained(repo_id,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLPipeline.from_pretrained(repo_id,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
else:
|
||||
if control_model_type != "stable-diffusion-xl-1.0-inpainting-0.1":
|
||||
controlnet = ControlNetModel.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
else:
|
||||
controlnet = ""
|
||||
).to("cuda")
|
||||
if function_choice == "img2img":
|
||||
model = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(repo_id,
|
||||
controlnet=controlnet,
|
||||
model = StableDiffusionXLControlNetImg2ImgPipeline.from_pipe(model,controlnet=controlnet,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
).to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLControlNetPipeline.from_pretrained(repo_id,
|
||||
controlnet=controlnet,
|
||||
torch_dtype=torch.float16,
|
||||
variant="fp16").to("cuda")
|
||||
model = StableDiffusionXLControlNetPipeline.from_pipe(
|
||||
model, controlnet=controlnet,
|
||||
torch_dtype=torch.float16).to("cuda")
|
||||
|
||||
else:
|
||||
|
||||
if function_choice == "img2img":
|
||||
model = StableDiffusionXLImg2ImgPipeline.from_pretrained(repo_id,
|
||||
torch_dtype=torch.float16,
|
||||
).to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLPipeline.from_pretrained(repo_id,
|
||||
torch_dtype=torch.float16,
|
||||
).to("cuda")
|
||||
|
||||
if control_model_type != "none":
|
||||
if not controlnet_repo_id:
|
||||
controlnet = ControlNetModel.from_unet(model.unet)
|
||||
cn_state_dict = load_file(controlnet_single_model, device="cuda")
|
||||
controlnet.load_state_dict(cn_state_dict, strict=False)
|
||||
|
||||
else:
|
||||
controlnet = ControlNetModel.from_pretrained(controlnet_repo_id, torch_dtype=torch.float16,
|
||||
).to("cuda")
|
||||
if function_choice == "img2img":
|
||||
model = StableDiffusionXLControlNetImg2ImgPipeline.from_pipe(model, controlnet=controlnet,
|
||||
torch_dtype=torch.float16,
|
||||
).to("cuda")
|
||||
else:
|
||||
model = StableDiffusionXLControlNetPipeline.from_pipe(
|
||||
model, controlnet=controlnet,
|
||||
torch_dtype=torch.float16).to("cuda")
|
||||
|
||||
if unet_model in sdxl_lightning_list:
|
||||
ckpt = folder_paths.get_full_path("unet", unet_model)
|
||||
unet = UNet2DConditionModel.from_config("stabilityai/stable-diffusion-xl-base-1.0",
|
||||
@@ -482,6 +460,11 @@ class HI_Diffusers_Model_Loader:
|
||||
model_info = ";".join([model_type, unet_model, control_model_type, function_choice,lora,trigger_words, adapter_info])
|
||||
torch.cuda.empty_cache()
|
||||
return (model, model_info,)
|
||||
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model, **kwargs):
|
||||
return remove_hidiffusion(model)
|
||||
|
||||
|
||||
class Hi_Sampler:
|
||||
@@ -656,19 +639,21 @@ class Hi_Sampler:
|
||||
image = kwargs["image"]
|
||||
image = input_size_adaptation_output(image, pre_input, width, height)
|
||||
if control_net == "stable-diffusion-xl-1.0-inpainting-0.1":
|
||||
print("controlnet inpainting")
|
||||
images = \
|
||||
model(prompt, negative_prompt=negative_prompt, image=image, mask_image=control_image,
|
||||
num_inference_steps=steps, guidance_scale=cfg, height=height, clip_skip=clip_skip,
|
||||
width=width, controlnet_conditioning_scale=controlnet_scale,
|
||||
seed=seed, ).images[0]
|
||||
else:
|
||||
print("controlnet img2img")
|
||||
images = model(prompt, negative_prompt=negative_prompt, image=image, control_image=control_image,
|
||||
num_inference_steps=steps, guidance_scale=cfg, height=height, width=width,
|
||||
clip_skip=clip_skip,
|
||||
controlnet_conditioning_scale=controlnet_scale,
|
||||
seed=seed, ).images[0]
|
||||
else:
|
||||
print("1234")
|
||||
print("controlnet txt2img")
|
||||
images = model(prompt, negative_prompt=negative_prompt, control_image=control_image,
|
||||
num_inference_steps=steps, guidance_scale=cfg, height=height, width=width,
|
||||
clip_skip=clip_skip,
|
||||
@@ -681,7 +666,7 @@ class Hi_Sampler:
|
||||
if adapter_info=="none":
|
||||
model.unfuse_lora()
|
||||
torch.cuda.empty_cache()
|
||||
del model
|
||||
|
||||
return (output_image,)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user