Merge branch 'master' of github.com:comfyanonymous/ComfyUI

This commit is contained in:
doctorpangloss 2024-06-03 11:42:55 -07:00
commit 3f559135c6
4 changed files with 46 additions and 19 deletions

View File

@ -14,6 +14,15 @@ import logging
MAX_PREVIEW_RESOLUTION = 512
def preview_to_image(latent_image):
latents_ubyte = (((latent_image + 1.0) / 2.0).clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
).to(device="cpu", dtype=torch.uint8, non_blocking=model_management.device_supports_non_blocking(latent_image.device))
return Image.fromarray(latents_ubyte.numpy())
class LatentPreviewer:
def decode_latent_to_preview(self, x0):
pass
@ -22,17 +31,14 @@ class LatentPreviewer:
preview_image = self.decode_latent_to_preview(x0)
return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION)
class TAESDPreviewerImpl(LatentPreviewer):
def __init__(self, taesd):
self.taesd = taesd
def decode_latent_to_preview(self, x0):
x_sample = self.taesd.decode(x0[:1])[0].detach()
x_sample = 255. * torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)
x_sample = np.moveaxis(x_sample.to(device="cpu", dtype=torch.uint8, non_blocking=model_management.device_supports_non_blocking(x_sample.device)).numpy(), 0, 2)
preview_image = Image.fromarray(x_sample)
return preview_image
x_sample = self.taesd.decode(x0[:1])[0].movedim(0, 2)
return preview_to_image(x_sample)
class Latent2RGBPreviewer(LatentPreviewer):
@ -42,13 +48,7 @@ class Latent2RGBPreviewer(LatentPreviewer):
def decode_latent_to_preview(self, x0):
self.latent_rgb_factors = self.latent_rgb_factors.to(dtype=x0.dtype, device=x0.device)
latent_image = x0[0].permute(1, 2, 0) @ self.latent_rgb_factors
latents_ubyte = (((latent_image + 1) / 2)
.clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
).to(device="cpu", dtype=torch.uint8, non_blocking=model_management.device_supports_non_blocking(latent_image.device))
return Image.fromarray(latents_ubyte.numpy())
return preview_to_image(latent_image)
def get_previewer(device, latent_format):
@ -88,6 +88,7 @@ def prepare_callback(model, steps, x0_output_dict=None):
previewer = get_previewer(model.load_device, model.model.latent_format)
pbar = utils.ProgressBar(steps)
def callback(step, x0, x, total_steps):
if x0_output_dict is not None:
x0_output_dict["x0"] = x0

View File

@ -964,6 +964,7 @@ def unload_all_models():
def resolve_lowvram_weight(weight, model, key): # TODO: remove
print("WARNING: The comfy.model_management.resolve_lowvram_weight function will be removed soon, please stop using it.")
return weight

View File

@ -88,9 +88,7 @@ class ModelPatcher(ModelManageable):
def model_size(self):
if self.size > 0:
return self.size
model_sd = self.model.state_dict()
self.size = model_management.module_size(self.model)
self.model_keys = set(model_sd.keys())
return self.size
def clone(self):
@ -102,7 +100,6 @@ class ModelPatcher(ModelManageable):
n.object_patches = self.object_patches.copy()
n.model_options = copy.deepcopy(self.model_options)
n.model_keys = self.model_keys
n.backup = self.backup
n.object_patches_backup = self.object_patches_backup
return n
@ -222,8 +219,9 @@ class ModelPatcher(ModelManageable):
def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
p = set()
model_sd = self.model.state_dict()
for k in patches:
if k in self.model_keys:
if k in model_sd:
p.add(k)
current_patches = self.patches.get(k, [])
current_patches.append((strength_patch, patches[k], strength_model))

View File

@ -59,12 +59,39 @@ class Guider_PerpNeg(samplers.CFGGuider):
self.neg_scale = neg_scale
def predict_noise(self, x, timestep, model_options={}, seed=None):
# in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
# but we'd rather do a single batch of sampling pos, neg, and empty, so we call calc_cond_batch([pos,neg,empty]) directly
positive_cond = self.conds.get("positive", None)
negative_cond = self.conds.get("negative", None)
empty_cond = self.conds.get("empty_negative_prompt", None)
out = samplers.calc_cond_batch(self.inner_model, [negative_cond, positive_cond, empty_cond], x, timestep, model_options)
return perp_neg(x, out[1], out[0], out[2], self.neg_scale, self.cfg)
(noise_pred_pos, noise_pred_neg, noise_pred_empty) = \
samplers.calc_cond_batch(self.inner_model, [positive_cond, negative_cond, empty_cond], x, timestep, model_options)
cfg_result = perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_empty, self.neg_scale, self.cfg)
# normally this would be done in cfg_function, but we skipped
# that for efficiency: we can compute the noise predictions in
# a single call to calc_cond_batch() (rather than two)
# so we replicate the hook here
for fn in model_options.get("sampler_post_cfg_function", []):
args = {
"denoised": cfg_result,
"cond": positive_cond,
"uncond": negative_cond,
"model": self.inner_model,
"uncond_denoised": noise_pred_neg,
"cond_denoised": noise_pred_pos,
"sigma": timestep,
"model_options": model_options,
"input": x,
# not in the original call in samplers.py:cfg_function, but made available for future hooks
"empty_cond": empty_cond,
"empty_cond_denoised": noise_pred_empty,}
# todo: is this supposed to be ** spread?
cfg_result = fn(args)
return cfg_result
class PerpNegGuider:
@classmethod