From bf3fc2f1b7f5b5cf684246be84838e6fc19aeb06 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 27 Sep 2023 16:45:22 -0400 Subject: [PATCH] Refactor sampling related code. --- comfy/sample.py | 24 ++++++++++++++---------- latent_preview.py | 18 ++++++++++++++++++ nodes.py | 17 +---------------- 3 files changed, 33 insertions(+), 26 deletions(-) diff --git a/comfy/sample.py b/comfy/sample.py index e4730b18..fe9f4118 100644 --- a/comfy/sample.py +++ b/comfy/sample.py @@ -70,25 +70,29 @@ def cleanup_additional_models(models): if hasattr(m, 'cleanup'): m.cleanup() -def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, 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): - device = comfy.model_management.get_torch_device() +def prepare_sampling(model, noise_shape, positive, negative, noise_mask): + device = model.load_device if noise_mask is not None: - noise_mask = prepare_mask(noise_mask, noise.shape, device) + noise_mask = prepare_mask(noise_mask, noise_shape, device) real_model = None models, inference_memory = get_additional_models(positive, negative, model.model_dtype()) - comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise.shape[0] * noise.shape[2] * noise.shape[3]) + inference_memory) + comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise_shape[0] * noise_shape[2] * noise_shape[3]) + inference_memory) real_model = model.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) + positive_copy = broadcast_cond(positive, noise_shape[0], device) + negative_copy = broadcast_cond(negative, noise_shape[0], device) + return real_model, positive_copy, negative_copy, noise_mask, models - sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options) +def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, 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): + real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask) + + noise = noise.to(model.load_device) + latent_image = latent_image.to(model.load_device) + + sampler = comfy.samplers.KSampler(real_model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options) samples = sampler.sample(noise, positive_copy, negative_copy, 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=callback, disable_pbar=disable_pbar, seed=seed) samples = samples.cpu() diff --git a/latent_preview.py b/latent_preview.py index 87240a58..740e0860 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -5,6 +5,7 @@ import numpy as np from comfy.cli_args import args, LatentPreviewMethod from comfy.taesd.taesd import TAESD import folder_paths +import comfy.utils MAX_PREVIEW_RESOLUTION = 512 @@ -74,4 +75,21 @@ def get_previewer(device, latent_format): previewer = Latent2RGBPreviewer(latent_format.latent_rgb_factors) return previewer +def prepare_callback(model, steps, x0_output_dict=None): + preview_format = "JPEG" + if preview_format not in ["JPEG", "PNG"]: + preview_format = "JPEG" + + previewer = get_previewer(model.load_device, model.model.latent_format) + + pbar = comfy.utils.ProgressBar(steps) + def callback(step, x0, x, total_steps): + if x0_output_dict is not None: + x0_output_dict["x0"] = x0 + + 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) + return callback diff --git a/nodes.py b/nodes.py index 4abb0d24..a847db6f 100644 --- a/nodes.py +++ b/nodes.py @@ -1189,11 +1189,8 @@ class SetLatentNoiseMask: s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) return (s,) - def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): - 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: @@ -1204,19 +1201,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, 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) - + callback = latent_preview.prepare_callback(model, steps) samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, seed=seed)