diff --git a/Hidiffusion_node.py b/Hidiffusion_node.py index 5bc6187..493b3c8 100644 --- a/Hidiffusion_node.py +++ b/Hidiffusion_node.py @@ -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,)