diff --git a/sampler/KSamplerWithRefiner.py b/sampler/KSamplerWithRefiner.py index d0c45fc..4884482 100644 --- a/sampler/KSamplerWithRefiner.py +++ b/sampler/KSamplerWithRefiner.py @@ -1,11 +1,18 @@ +import os import torch import comfy.model_management +import comfy.utils import math import numpy as np import latent_preview from comfy.sample import * from . import samplers_advanced +from .adm_patch import patch_negative_adm + + +patch_negative_adm() + def sample_refined(model, refiner_model, noise, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None): """ @@ -82,6 +89,85 @@ def ksampler_refined(model, refiner_model, seed, steps, cfg, sampler_name, sched out["samples"] = samples return (out, ) + +@torch.no_grad() +def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive, refiner_negative, latent, + seed=None, steps=30, refiner_switch_step=20, cfg=7.0, sampler_name='dpmpp_2m_sde_gpu', + scheduler='karras', denoise=1.0, disable_noise=False, start_step=None, last_step=None, + force_full_denoise=False, callback_function=None): + # SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"] + # SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral", + # "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", + # "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2"] + + seed = seed if isinstance(seed, int) else random.randint(1, 2 ** 64) + + device = comfy.model_management.get_torch_device() + latent_image = latent["samples"] + + if disable_noise: + noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") + else: + batch_inds = latent["batch_index"] if "batch_index" in latent else None + noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds) + + noise_mask = None + if "noise_mask" in latent: + noise_mask = latent["noise_mask"] + + preview_format = "JPEG" + if preview_format not in ["JPEG", "PNG"]: + preview_format = "JPEG" + + previewer = latent_preview.get_previewer(device, model.model.latent_format) + + pbar = comfy.utils.ProgressBar(steps) + def callback(step, x0, x, total_steps): + preview_bytes = None + if previewer: + preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0) + pbar.update_absolute(step + 1, total_steps, preview_bytes) + + sigmas = None + disable_pbar = False + + if noise_mask is not None: + noise_mask = prepare_mask(noise_mask, noise.shape, device) + + comfy.model_management.load_model_gpu(model) + + noise = noise.to(device) + latent_image = latent_image.to(device) + + positive_copy = broadcast_cond(positive, noise.shape[0], device) + negative_copy = broadcast_cond(negative, noise.shape[0], device) + + refiner_positive_copy = broadcast_cond(refiner_positive, noise.shape[0], device) + refiner_negative_copy = broadcast_cond(refiner_negative, noise.shape[0], device) + + models = load_additional_models(positive, negative, model.model_dtype()) + + sampler = samplers_advanced.KSamplerWithRefiner(model=model, refiner_model=refiner, steps=steps, device=device, + sampler=sampler_name, scheduler=scheduler, + denoise=denoise, model_options=model.model_options) + + samples = sampler.sample(noise, positive_copy, negative_copy, refiner_positive=refiner_positive_copy, + refiner_negative=refiner_negative_copy, refiner_switch_step=refiner_switch_step, + cfg=cfg, latent_image=latent_image, + start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, + denoise_mask=noise_mask, sigmas=sigmas, callback_function=callback, disable_pbar=disable_pbar, + seed=seed) + + samples = samples.cpu() + + cleanup_additional_models(models) + + out = latent.copy() + out["samples"] = samples + + return out + + class KSamplerWithRefiner: @classmethod def INPUT_TYPES(s): @@ -90,10 +176,11 @@ class KSamplerWithRefiner: "refiner_model": ("MODEL",), "add_noise": (["enable", "disable"], ), "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}), - "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), - "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), + "steps": ("INT", {"default": 30, "min": 1, "max": 10000}), + "refiner_switch_step": ("INT", {"default": 20, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "dpmpp_2m_sde_gpu", }), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras", }), "positive": ("CONDITIONING", ), "negative": ("CONDITIONING", ), "refiner_positive": ("CONDITIONING", ), @@ -110,11 +197,12 @@ class KSamplerWithRefiner: CATEGORY = "sampling" - def sample(self, model, refiner_model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0): + def sample(self, model, refiner_model, add_noise, noise_seed, steps, refiner_switch_step, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0): force_full_denoise = True if return_with_leftover_noise == "enable": force_full_denoise = False disable_noise = False if add_noise == "disable": disable_noise = True - return ksampler_refined(model, refiner_model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) + # return ksampler_refined(model, refiner_model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, refiner_positive, refiner_negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) + return ksampler_with_refiner(model, positive, negative, refiner_model, refiner_positive, refiner_negative, latent_image, noise_seed, steps, refiner_switch_step, cfg, sampler_name, scheduler, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) diff --git a/sampler/adm_patch.py b/sampler/adm_patch.py new file mode 100644 index 0000000..7567fa9 --- /dev/null +++ b/sampler/adm_patch.py @@ -0,0 +1,33 @@ +import torch +import comfy.model_base + + +def sdxl_encode_adm_patched(self, **kwargs): + clip_pooled = kwargs["pooled_output"] + width = kwargs.get("width", 768) + height = kwargs.get("height", 768) + crop_w = kwargs.get("crop_w", 0) + crop_h = kwargs.get("crop_h", 0) + target_width = kwargs.get("target_width", width) + target_height = kwargs.get("target_height", height) + + if kwargs.get("prompt_type", "") == "negative": + width *= 0.8 + height *= 0.8 + elif kwargs.get("prompt_type", "") == "positive": + width *= 1.5 + height *= 1.5 + + out = [] + out.append(self.embedder(torch.Tensor([height]))) + out.append(self.embedder(torch.Tensor([width]))) + out.append(self.embedder(torch.Tensor([crop_h]))) + out.append(self.embedder(torch.Tensor([crop_w]))) + out.append(self.embedder(torch.Tensor([target_height]))) + out.append(self.embedder(torch.Tensor([target_width]))) + flat = torch.flatten(torch.cat(out))[None, ] + return torch.cat((clip_pooled.to(flat.device), flat), dim=1) + + +def patch_negative_adm(): + comfy.model_base.SDXL.encode_adm = sdxl_encode_adm_patched diff --git a/sampler/cv2win32.py b/sampler/cv2win32.py new file mode 100644 index 0000000..5e41dbc --- /dev/null +++ b/sampler/cv2win32.py @@ -0,0 +1,35 @@ +import threading +import cv2 + + +buffer = [] + + +def worker(): + global buffer + while True: + cv2.waitKey(50) + try: + if len(buffer) > 0: + task = buffer.pop(0) + if task is None: + cv2.destroyAllWindows() + else: + flag, img, title = task + cv2.imshow(flag, img) + cv2.setWindowTitle(flag, title) + cv2.setWindowProperty(flag, cv2.WND_PROP_TOPMOST, 1) + except Exception as e: + print(e) + pass + + +def show_preview(flag, img, title='preview'): + buffer.append((flag, img[..., ::-1].copy(), title)) + + +def close_all_preview(): + buffer.append(None) + + +threading.Thread(target=worker, daemon=True).start() diff --git a/sampler/samplers_advanced.py b/sampler/samplers_advanced.py index 183d8c4..e68c059 100644 --- a/sampler/samplers_advanced.py +++ b/sampler/samplers_advanced.py @@ -19,8 +19,8 @@ class KSamplerWithRefiner: self.model_patcher = model self.refiner_model_patcher = refiner_model - self.model = model - self.refiner_model = refiner_model + self.model = model.model + self.refiner_model = refiner_model.model self.model_denoise = CFGNoisePredictor(self.model) self.refiner_model_denoise = CFGNoisePredictor(self.refiner_model) diff --git a/sampler/test ksampler refiner.json b/sampler/test ksampler refiner.json index 279a235..7e91a90 100644 --- a/sampler/test ksampler refiner.json +++ b/sampler/test ksampler refiner.json @@ -271,32 +271,6 @@ "sdxl_vae.safetensors" ] }, - { - "id": 28, - "type": "SaveImage", - "pos": [ - 2465, - 298 - ], - "size": { - "0": 315, - "1": 58 - }, - "flags": {}, - "order": 17, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 51 - } - ], - "properties": {}, - "widgets_values": [ - "ComfyUI" - ] - }, { "id": 17, "type": "CLIPTextEncodeSDXL", @@ -584,7 +558,7 @@ ], "size": { "0": 315, - "1": 334 + "1": 546 }, "flags": {}, "order": 11, @@ -629,12 +603,12 @@ "enable", 1, "fixed", - 20, + 30, 8, - "dpmpp_2m", + "dpmpp_2m_sde_gpu", "karras", 0, - 16, + 20, "enable" ] }, @@ -707,9 +681,10 @@ "enable", 1, "fixed", + 30, 20, - 8, - "dpmpp_2m", + 7, + "dpmpp_2m_sde_gpu", "karras", 0, 10000, @@ -725,7 +700,7 @@ ], "size": { "0": 315, - "1": 334 + "1": 546 }, "flags": {}, "order": 13, @@ -770,11 +745,11 @@ "disable", 1, "fixed", - 20, + 30, 8, - "dpmpp_2m", + "dpmpp_2m_sde_gpu", "karras", - 16, + 20, 10000, "disable" ] @@ -819,32 +794,6 @@ "Node name for S&R": "VAEDecode" } }, - { - "id": 9, - "type": "SaveImage", - "pos": [ - 2014.1724880000002, - -373.70633999999995 - ], - "size": { - "0": 210, - "1": 58 - }, - "flags": {}, - "order": 16, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 9 - } - ], - "properties": {}, - "widgets_values": [ - "ComfyUI" - ] - }, { "id": 27, "type": "VAEDecode", @@ -885,6 +834,58 @@ "properties": { "Node name for S&R": "VAEDecode" } + }, + { + "id": 9, + "type": "SaveImage", + "pos": [ + 2014.1724880000002, + -373.70633999999995 + ], + "size": { + "0": 210, + "1": 270 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 28, + "type": "SaveImage", + "pos": [ + 2465, + 298 + ], + "size": { + "0": 315, + "1": 270 + }, + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 51 + } + ], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] } ], "links": [