ComfyUI version v0.3.41

Include alpha in sampling and minor refactoring

.github/workflows/test-unit.yml

@@ -28,7 +28,3 @@ jobs:
       run: |
         pip install -r tests-unit/requirements.txt
         python -m pytest tests-unit
-    - name: Run Execution Model Tests
-      run: |
-        python -m pytest tests/inference/test_execution.py
-

comfy/k_diffusion/sampling.py

@@ -1,4 +1,5 @@
 import math
+from functools import partial
 
 from scipy import integrate
 import torch
@@ -142,6 +143,33 @@ class BrownianTreeNoiseSampler:
         return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
 
 
+def sigma_to_half_log_snr(sigma, model_sampling):
+    """Convert sigma to half-logSNR log(alpha_t / sigma_t)."""
+    if isinstance(model_sampling, comfy.model_sampling.CONST):
+        # log((1 - t) / t) = log((1 - sigma) / sigma)
+        return sigma.logit().neg()
+    return sigma.log().neg()
+
+
+def half_log_snr_to_sigma(half_log_snr, model_sampling):
+    """Convert half-logSNR log(alpha_t / sigma_t) to sigma."""
+    if isinstance(model_sampling, comfy.model_sampling.CONST):
+        # 1 / (1 + exp(half_log_snr))
+        return half_log_snr.neg().sigmoid()
+    return half_log_snr.neg().exp()
+
+
+def offset_first_sigma_for_snr(sigmas, model_sampling, percent_offset=1e-4):
+    """Adjust the first sigma to avoid invalid logSNR."""
+    if len(sigmas) <= 1:
+        return sigmas
+    if isinstance(model_sampling, comfy.model_sampling.CONST):
+        if sigmas[0] >= 1:
+            sigmas = sigmas.clone()
+            sigmas[0] = model_sampling.percent_to_sigma(percent_offset)
+    return sigmas
+
+
 @torch.no_grad()
 def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
     """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
@@ -753,6 +781,7 @@ def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=No
         old_denoised = denoised
     return x
 
+
 @torch.no_grad()
 def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
     """DPM-Solver++(2M) SDE."""
@@ -768,9 +797,12 @@ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
     s_in = x.new_ones([x.shape[0]])
 
+    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
+    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
+    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
+
     old_denoised = None
-    h_last = None
-    h = None
+    h, h_last = None, None
 
     for i in trange(len(sigmas) - 1, disable=disable):
         denoised = model(x, sigmas[i] * s_in, **extra_args)
@@ -781,26 +813,29 @@ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
             x = denoised
         else:
             # DPM-Solver++(2M) SDE
-            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
-            h = s - t
-            eta_h = eta * h
+            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+            h = lambda_t - lambda_s
+            h_eta = h * (eta + 1)
 
-            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised
+            alpha_t = sigmas[i + 1] * lambda_t.exp()
+
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x + alpha_t * (-h_eta).expm1().neg() * denoised
 
             if old_denoised is not None:
                 r = h_last / h
                 if solver_type == 'heun':
-                    x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised)
+                    x = x + alpha_t * ((-h_eta).expm1().neg() / (-h_eta) + 1) * (1 / r) * (denoised - old_denoised)
                 elif solver_type == 'midpoint':
-                    x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)
+                    x = x + 0.5 * alpha_t * (-h_eta).expm1().neg() * (1 / r) * (denoised - old_denoised)
 
-            if eta:
-                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
+            if eta > 0 and s_noise > 0:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
 
         old_denoised = denoised
         h_last = h
     return x
 
+
 @torch.no_grad()
 def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
     """DPM-Solver++(3M) SDE."""
@@ -814,6 +849,10 @@ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
     s_in = x.new_ones([x.shape[0]])
 
+    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
+    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
+    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
+
     denoised_1, denoised_2 = None, None
     h, h_1, h_2 = None, None, None
 
@@ -825,13 +864,16 @@ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
             # Denoising step
             x = denoised
         else:
-            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
-            h = s - t
+            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+            h = lambda_t - lambda_s
             h_eta = h * (eta + 1)
 
-            x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised
+            alpha_t = sigmas[i + 1] * lambda_t.exp()
+
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x + alpha_t * (-h_eta).expm1().neg() * denoised
 
             if h_2 is not None:
+                # DPM-Solver++(3M) SDE
                 r0 = h_1 / h
                 r1 = h_2 / h
                 d1_0 = (denoised - denoised_1) / r0
@@ -840,20 +882,22 @@ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
                 d2 = (d1_0 - d1_1) / (r0 + r1)
                 phi_2 = h_eta.neg().expm1() / h_eta + 1
                 phi_3 = phi_2 / h_eta - 0.5
-                x = x + phi_2 * d1 - phi_3 * d2
+                x = x + (alpha_t * phi_2) * d1 - (alpha_t * phi_3) * d2
             elif h_1 is not None:
+                # DPM-Solver++(2M) SDE
                 r = h_1 / h
                 d = (denoised - denoised_1) / r
                 phi_2 = h_eta.neg().expm1() / h_eta + 1
-                x = x + phi_2 * d
+                x = x + (alpha_t * phi_2) * d
 
-            if eta:
+            if eta > 0 and s_noise > 0:
                 x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
 
         denoised_1, denoised_2 = denoised, denoised_1
         h_1, h_2 = h, h_1
     return x
 
+
 @torch.no_grad()
 def sample_dpmpp_3m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
     if len(sigmas) <= 1:
@@ -863,6 +907,7 @@ def sample_dpmpp_3m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, di
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
     return sample_dpmpp_3m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
 
+
 @torch.no_grad()
 def sample_dpmpp_2m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
     if len(sigmas) <= 1:
@@ -872,6 +917,7 @@ def sample_dpmpp_2m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, di
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
     return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
 
+
 @torch.no_grad()
 def sample_dpmpp_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
     if len(sigmas) <= 1:
@@ -1449,12 +1495,12 @@ def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None
         old_denoised = denoised
     return x
 
+
 @torch.no_grad()
 def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5):
-    '''
-    SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VE Data Prediction) stage 2
-    Arxiv: https://arxiv.org/abs/2305.14267
-    '''
+    """SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 2.
+    arXiv: https://arxiv.org/abs/2305.14267
+    """
     extra_args = {} if extra_args is None else extra_args
     seed = extra_args.get("seed", None)
     noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
@@ -1462,6 +1508,11 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
 
     inject_noise = eta > 0 and s_noise > 0
 
+    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
+    sigma_fn = partial(half_log_snr_to_sigma, model_sampling=model_sampling)
+    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
+    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
+
     for i in trange(len(sigmas) - 1, disable=disable):
         denoised = model(x, sigmas[i] * s_in, **extra_args)
         if callback is not None:
@@ -1469,80 +1520,96 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
         if sigmas[i + 1] == 0:
             x = denoised
         else:
-            t, t_next = -sigmas[i].log(), -sigmas[i + 1].log()
-            h = t_next - t
+            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+            h = lambda_t - lambda_s
             h_eta = h * (eta + 1)
-            s = t + r * h
+            lambda_s_1 = lambda_s + r * h
             fac = 1 / (2 * r)
-            sigma_s = s.neg().exp()
+            sigma_s_1 = sigma_fn(lambda_s_1)
+
+            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
+            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+            alpha_t = sigmas[i + 1] * lambda_t.exp()
 
             coeff_1, coeff_2 = (-r * h_eta).expm1(), (-h_eta).expm1()
             if inject_noise:
+                # 0 < r < 1
                 noise_coeff_1 = (-2 * r * h * eta).expm1().neg().sqrt()
-                noise_coeff_2 = ((-2 * r * h * eta).expm1() - (-2 * h * eta).expm1()).sqrt()
-                noise_1, noise_2 = noise_sampler(sigmas[i], sigma_s), noise_sampler(sigma_s, sigmas[i + 1])
+                noise_coeff_2 = (-r * h * eta).exp() * (-2 * (1 - r) * h * eta).expm1().neg().sqrt()
+                noise_1, noise_2 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigmas[i + 1])
 
             # Step 1
-            x_2 = (coeff_1 + 1) * x - coeff_1 * denoised
-            if inject_noise:
-                x_2 = x_2 + sigma_s * (noise_coeff_1 * noise_1) * s_noise
-            denoised_2 = model(x_2, sigma_s * s_in, **extra_args)
-
-            # Step 2
-            denoised_d = (1 - fac) * denoised + fac * denoised_2
-            x = (coeff_2 + 1) * x - coeff_2 * denoised_d
-            if inject_noise:
-                x = x + sigmas[i + 1] * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
-    return x
-
-@torch.no_grad()
-def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
-    '''
-    SEEDS-3 - Stochastic Explicit Exponential Derivative-free Solvers (VE Data Prediction) stage 3
-    Arxiv: https://arxiv.org/abs/2305.14267
-    '''
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    inject_noise = eta > 0 and s_noise > 0
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            x = denoised
-        else:
-            t, t_next = -sigmas[i].log(), -sigmas[i + 1].log()
-            h = t_next - t
-            h_eta = h * (eta + 1)
-            s_1 = t + r_1 * h
-            s_2 = t + r_2 * h
-            sigma_s_1, sigma_s_2 = s_1.neg().exp(), s_2.neg().exp()
-
-            coeff_1, coeff_2, coeff_3 = (-r_1 * h_eta).expm1(), (-r_2 * h_eta).expm1(), (-h_eta).expm1()
-            if inject_noise:
-                noise_coeff_1 = (-2 * r_1 * h * eta).expm1().neg().sqrt()
-                noise_coeff_2 = ((-2 * r_1 * h * eta).expm1() - (-2 * r_2 * h * eta).expm1()).sqrt()
-                noise_coeff_3 = ((-2 * r_2 * h * eta).expm1() - (-2 * h * eta).expm1()).sqrt()
-                noise_1, noise_2, noise_3 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigma_s_2), noise_sampler(sigma_s_2, sigmas[i + 1])
-
-            # Step 1
-            x_2 = (coeff_1 + 1) * x - coeff_1 * denoised
+            x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
             if inject_noise:
                 x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
             denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
 
             # Step 2
-            x_3 = (coeff_2 + 1) * x - coeff_2 * denoised + (r_2 / r_1) * (coeff_2 / (r_2 * h_eta) + 1) * (denoised_2 - denoised)
+            denoised_d = (1 - fac) * denoised + fac * denoised_2
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_2 * denoised_d
+            if inject_noise:
+                x = x + sigmas[i + 1] * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
+    return x
+
+
+@torch.no_grad()
+def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
+    """SEEDS-3 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 3.
+    arXiv: https://arxiv.org/abs/2305.14267
+    """
+    extra_args = {} if extra_args is None else extra_args
+    seed = extra_args.get("seed", None)
+    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+
+    inject_noise = eta > 0 and s_noise > 0
+
+    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
+    sigma_fn = partial(half_log_snr_to_sigma, model_sampling=model_sampling)
+    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
+    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            x = denoised
+        else:
+            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+            h = lambda_t - lambda_s
+            h_eta = h * (eta + 1)
+            lambda_s_1 = lambda_s + r_1 * h
+            lambda_s_2 = lambda_s + r_2 * h
+            sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
+
+            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
+            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+            alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
+            alpha_t = sigmas[i + 1] * lambda_t.exp()
+
+            coeff_1, coeff_2, coeff_3 = (-r_1 * h_eta).expm1(), (-r_2 * h_eta).expm1(), (-h_eta).expm1()
+            if inject_noise:
+                # 0 < r_1 < r_2 < 1
+                noise_coeff_1 = (-2 * r_1 * h * eta).expm1().neg().sqrt()
+                noise_coeff_2 = (-r_1 * h * eta).exp() * (-2 * (r_2 - r_1) * h * eta).expm1().neg().sqrt()
+                noise_coeff_3 = (-r_2 * h * eta).exp() * (-2 * (1 - r_2) * h * eta).expm1().neg().sqrt()
+                noise_1, noise_2, noise_3 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigma_s_2), noise_sampler(sigma_s_2, sigmas[i + 1])
+
+            # Step 1
+            x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
+            if inject_noise:
+                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
+            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+
+            # Step 2
+            x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * coeff_2 * denoised + (r_2 / r_1) * alpha_s_2 * (coeff_2 / (r_2 * h_eta) + 1) * (denoised_2 - denoised)
             if inject_noise:
                 x_3 = x_3 + sigma_s_2 * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
             denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
 
             # Step 3
-            x = (coeff_3 + 1) * x - coeff_3 * denoised + (1. / r_2) * (coeff_3 / h_eta + 1) * (denoised_3 - denoised)
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_3 * denoised + (1. / r_2) * alpha_t * (coeff_3 / h_eta + 1) * (denoised_3 - denoised)
             if inject_noise:
                 x = x + sigmas[i + 1] * (noise_coeff_3 * noise_1 + noise_coeff_2 * noise_2 + noise_coeff_1 * noise_3) * s_noise
     return x

comfy/ldm/cosmos/predict2.py

@@ -70,11 +70,7 @@ def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H
     q_B_H_S_D = rearrange(q_B_S_H_D, "b ... h k -> b h ... k").view(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1])
     k_B_H_S_D = rearrange(k_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
     v_B_H_S_D = rearrange(v_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
-    result_B_S_HD = rearrange(
-        optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True, skip_output_reshape=True), "b h ... l -> b ... (h l)"
-    )
-
-    return result_B_S_HD
+    return optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True)
 
 
 class Attention(nn.Module):

comfy/ldm/lightricks/model.py

@@ -261,8 +261,8 @@ class CrossAttention(nn.Module):
         self.heads = heads
         self.dim_head = dim_head
 
-        self.q_norm = operations.RMSNorm(inner_dim, dtype=dtype, device=device)
-        self.k_norm = operations.RMSNorm(inner_dim, dtype=dtype, device=device)
+        self.q_norm = operations.RMSNorm(inner_dim, eps=1e-5, dtype=dtype, device=device)
+        self.k_norm = operations.RMSNorm(inner_dim, eps=1e-5, dtype=dtype, device=device)
 
         self.to_q = operations.Linear(query_dim, inner_dim, bias=True, dtype=dtype, device=device)
         self.to_k = operations.Linear(context_dim, inner_dim, bias=True, dtype=dtype, device=device)

comfy/model_base.py

@@ -1014,9 +1014,30 @@ class CosmosPredict2(BaseModel):
         if cross_attn is not None:
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
 
+        denoise_mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None))
+        if denoise_mask is not None:
+            out["denoise_mask"] = comfy.conds.CONDRegular(denoise_mask)
+
         out['fps'] = comfy.conds.CONDConstant(kwargs.get("frame_rate", None))
         return out
 
+    def process_timestep(self, timestep, x, denoise_mask=None, **kwargs):
+        if denoise_mask is None:
+            return timestep
+        condition_video_mask_B_1_T_1_1 = denoise_mask.mean(dim=[1, 3, 4], keepdim=True)
+        c_noise_B_1_T_1_1 = 0.0 * (1.0 - condition_video_mask_B_1_T_1_1) + timestep.reshape(timestep.shape[0], 1, 1, 1, 1) * condition_video_mask_B_1_T_1_1
+        out = c_noise_B_1_T_1_1.squeeze(dim=[1, 3, 4])
+        return out
+
+    def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
+        sigma = sigma.reshape([sigma.shape[0]] + [1] * (len(noise.shape) - 1))
+        sigma_noise_augmentation = 0 #TODO
+        if sigma_noise_augmentation != 0:
+            latent_image = latent_image + noise
+        latent_image = self.model_sampling.calculate_input(torch.tensor([sigma_noise_augmentation], device=latent_image.device, dtype=latent_image.dtype), latent_image)
+        sigma = (sigma / (sigma + 1))
+        return latent_image / (1.0 - sigma)
+
 class Lumina2(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.lumina.model.NextDiT)

comfy/model_detection.py

@@ -441,11 +441,16 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
             dit_config["rope_h_extrapolation_ratio"] = 4.0
             dit_config["rope_w_extrapolation_ratio"] = 4.0
             dit_config["rope_t_extrapolation_ratio"] = 1.0
-        elif dit_config["in_channels"] == 17:
-            dit_config["extra_per_block_abs_pos_emb"] = False
-            dit_config["rope_h_extrapolation_ratio"] = 3.0
-            dit_config["rope_w_extrapolation_ratio"] = 3.0
-            dit_config["rope_t_extrapolation_ratio"] = 1.0
+        elif dit_config["in_channels"] == 17: # img to video
+            if dit_config["model_channels"] == 2048:
+                dit_config["extra_per_block_abs_pos_emb"] = False
+                dit_config["rope_h_extrapolation_ratio"] = 3.0
+                dit_config["rope_w_extrapolation_ratio"] = 3.0
+                dit_config["rope_t_extrapolation_ratio"] = 1.0
+            elif dit_config["model_channels"] == 5120:
+                dit_config["rope_h_extrapolation_ratio"] = 2.0
+                dit_config["rope_w_extrapolation_ratio"] = 2.0
+                dit_config["rope_t_extrapolation_ratio"] = 0.8333333333333334
 
         dit_config["extra_h_extrapolation_ratio"] = 1.0
         dit_config["extra_w_extrapolation_ratio"] = 1.0

comfy/sd1_clip.py

@@ -462,7 +462,7 @@ class SDTokenizer:
             tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
         self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path, **tokenizer_args)
         self.max_length = tokenizer_data.get("{}_max_length".format(embedding_key), max_length)
-        self.min_length = min_length
+        self.min_length = tokenizer_data.get("{}_min_length".format(embedding_key), min_length)
         self.end_token = None
         self.min_padding = min_padding
 

comfy/utils.py

@@ -997,12 +997,11 @@ def set_progress_bar_global_hook(function):
     PROGRESS_BAR_HOOK = function
 
 class ProgressBar:
-    def __init__(self, total, node_id=None):
+    def __init__(self, total):
         global PROGRESS_BAR_HOOK
         self.total = total
         self.current = 0
         self.hook = PROGRESS_BAR_HOOK
-        self.node_id = node_id
 
     def update_absolute(self, value, total=None, preview=None):
         if total is not None:
@@ -1011,7 +1010,7 @@ class ProgressBar:
             value = self.total
         self.current = value
         if self.hook is not None:
-            self.hook(self.current, self.total, preview, node_id=self.node_id)
+            self.hook(self.current, self.total, preview)
 
     def update(self, value):
         self.update_absolute(self.current + value)

comfy_execution/caching.py

@@ -1,7 +1,6 @@
 import itertools
 from typing import Sequence, Mapping, Dict
 from comfy_execution.graph import DynamicPrompt
-from abc import ABC, abstractmethod
 
 import nodes
 
@@ -17,13 +16,12 @@ def include_unique_id_in_input(class_type: str) -> bool:
     NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] = "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values()
     return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type]
 
-class CacheKeySet(ABC):
+class CacheKeySet:
     def __init__(self, dynprompt, node_ids, is_changed_cache):
         self.keys = {}
         self.subcache_keys = {}
 
-    @abstractmethod
-    async def add_keys(self, node_ids):
+    def add_keys(self, node_ids):
         raise NotImplementedError()
 
     def all_node_ids(self):
@@ -62,8 +60,9 @@ class CacheKeySetID(CacheKeySet):
     def __init__(self, dynprompt, node_ids, is_changed_cache):
         super().__init__(dynprompt, node_ids, is_changed_cache)
         self.dynprompt = dynprompt
+        self.add_keys(node_ids)
 
-    async def add_keys(self, node_ids):
+    def add_keys(self, node_ids):
         for node_id in node_ids:
             if node_id in self.keys:
                 continue
@@ -78,36 +77,37 @@ class CacheKeySetInputSignature(CacheKeySet):
         super().__init__(dynprompt, node_ids, is_changed_cache)
         self.dynprompt = dynprompt
         self.is_changed_cache = is_changed_cache
+        self.add_keys(node_ids)
 
     def include_node_id_in_input(self) -> bool:
         return False
 
-    async def add_keys(self, node_ids):
+    def add_keys(self, node_ids):
         for node_id in node_ids:
             if node_id in self.keys:
                 continue
             if not self.dynprompt.has_node(node_id):
                 continue
             node = self.dynprompt.get_node(node_id)
-            self.keys[node_id] = await self.get_node_signature(self.dynprompt, node_id)
+            self.keys[node_id] = self.get_node_signature(self.dynprompt, node_id)
             self.subcache_keys[node_id] = (node_id, node["class_type"])
 
-    async def get_node_signature(self, dynprompt, node_id):
+    def get_node_signature(self, dynprompt, node_id):
         signature = []
         ancestors, order_mapping = self.get_ordered_ancestry(dynprompt, node_id)
-        signature.append(await self.get_immediate_node_signature(dynprompt, node_id, order_mapping))
+        signature.append(self.get_immediate_node_signature(dynprompt, node_id, order_mapping))
         for ancestor_id in ancestors:
-            signature.append(await self.get_immediate_node_signature(dynprompt, ancestor_id, order_mapping))
+            signature.append(self.get_immediate_node_signature(dynprompt, ancestor_id, order_mapping))
         return to_hashable(signature)
 
-    async def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
+    def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
         if not dynprompt.has_node(node_id):
             # This node doesn't exist -- we can't cache it.
             return [float("NaN")]
         node = dynprompt.get_node(node_id)
         class_type = node["class_type"]
         class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-        signature = [class_type, await self.is_changed_cache.get(node_id)]
+        signature = [class_type, self.is_changed_cache.get(node_id)]
         if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or include_unique_id_in_input(class_type):
             signature.append(node_id)
         inputs = node["inputs"]
@@ -150,10 +150,9 @@ class BasicCache:
         self.cache = {}
         self.subcaches = {}
 
-    async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
+    def set_prompt(self, dynprompt, node_ids, is_changed_cache):
         self.dynprompt = dynprompt
         self.cache_key_set = self.key_class(dynprompt, node_ids, is_changed_cache)
-        await self.cache_key_set.add_keys(node_ids)
         self.is_changed_cache = is_changed_cache
         self.initialized = True
 
@@ -202,13 +201,13 @@ class BasicCache:
         else:
             return None
 
-    async def _ensure_subcache(self, node_id, children_ids):
+    def _ensure_subcache(self, node_id, children_ids):
         subcache_key = self.cache_key_set.get_subcache_key(node_id)
         subcache = self.subcaches.get(subcache_key, None)
         if subcache is None:
             subcache = BasicCache(self.key_class)
             self.subcaches[subcache_key] = subcache
-        await subcache.set_prompt(self.dynprompt, children_ids, self.is_changed_cache)
+        subcache.set_prompt(self.dynprompt, children_ids, self.is_changed_cache)
         return subcache
 
     def _get_subcache(self, node_id):
@@ -260,10 +259,10 @@ class HierarchicalCache(BasicCache):
         assert cache is not None
         cache._set_immediate(node_id, value)
 
-    async def ensure_subcache_for(self, node_id, children_ids):
+    def ensure_subcache_for(self, node_id, children_ids):
         cache = self._get_cache_for(node_id)
         assert cache is not None
-        return await cache._ensure_subcache(node_id, children_ids)
+        return cache._ensure_subcache(node_id, children_ids)
 
 class LRUCache(BasicCache):
     def __init__(self, key_class, max_size=100):
@@ -274,8 +273,8 @@ class LRUCache(BasicCache):
         self.used_generation = {}
         self.children = {}
 
-    async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
-        await super().set_prompt(dynprompt, node_ids, is_changed_cache)
+    def set_prompt(self, dynprompt, node_ids, is_changed_cache):
+        super().set_prompt(dynprompt, node_ids, is_changed_cache)
         self.generation += 1
         for node_id in node_ids:
             self._mark_used(node_id)
@@ -304,11 +303,11 @@ class LRUCache(BasicCache):
         self._mark_used(node_id)
         return self._set_immediate(node_id, value)
 
-    async def ensure_subcache_for(self, node_id, children_ids):
+    def ensure_subcache_for(self, node_id, children_ids):
         # Just uses subcaches for tracking 'live' nodes
-        await super()._ensure_subcache(node_id, children_ids)
+        super()._ensure_subcache(node_id, children_ids)
 
-        await self.cache_key_set.add_keys(children_ids)
+        self.cache_key_set.add_keys(children_ids)
         self._mark_used(node_id)
         cache_key = self.cache_key_set.get_data_key(node_id)
         self.children[cache_key] = []
@@ -338,7 +337,7 @@ class DependencyAwareCache(BasicCache):
         self.ancestors = {}    # Maps node_id -> set of ancestor node_ids
         self.executed_nodes = set()  # Tracks nodes that have been executed
 
-    async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
+    def set_prompt(self, dynprompt, node_ids, is_changed_cache):
         """
         Clear the entire cache and rebuild the dependency graph.
 
@@ -355,7 +354,7 @@ class DependencyAwareCache(BasicCache):
         self.executed_nodes.clear()
 
         # Call the parent method to initialize the cache with the new prompt
-        await super().set_prompt(dynprompt, node_ids, is_changed_cache)
+        super().set_prompt(dynprompt, node_ids, is_changed_cache)
 
         # Rebuild the dependency graph
         self._build_dependency_graph(dynprompt, node_ids)
@@ -406,7 +405,7 @@ class DependencyAwareCache(BasicCache):
         """
         return self._get_immediate(node_id)
 
-    async def ensure_subcache_for(self, node_id, children_ids):
+    def ensure_subcache_for(self, node_id, children_ids):
         """
         Ensure a subcache exists for a node and update dependencies.
 
@@ -417,7 +416,7 @@ class DependencyAwareCache(BasicCache):
         Returns:
             The subcache object for the node.
         """
-        subcache = await super()._ensure_subcache(node_id, children_ids)
+        subcache = super()._ensure_subcache(node_id, children_ids)
         for child_id in children_ids:
             self.descendants[node_id].add(child_id)
             self.ancestors[child_id].add(node_id)

comfy_execution/graph.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 from typing import Type, Literal
 
 import nodes
-import asyncio
 from comfy_execution.graph_utils import is_link
 from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, InputTypeOptions
 
@@ -101,8 +100,6 @@ class TopologicalSort:
         self.pendingNodes = {}
         self.blockCount = {} # Number of nodes this node is directly blocked by
         self.blocking = {} # Which nodes are blocked by this node
-        self.externalBlocks = 0
-        self.unblockedEvent = asyncio.Event()
 
     def get_input_info(self, unique_id, input_name):
         class_type = self.dynprompt.get_node(unique_id)["class_type"]
@@ -156,16 +153,6 @@ class TopologicalSort:
         for link in links:
             self.add_strong_link(*link)
 
-    def add_external_block(self, node_id):
-        assert node_id in self.blockCount, "Can't add external block to a node that isn't pending"
-        self.externalBlocks += 1
-        self.blockCount[node_id] += 1
-        def unblock():
-            self.externalBlocks -= 1
-            self.blockCount[node_id] -= 1
-            self.unblockedEvent.set()
-        return unblock
-
     def is_cached(self, node_id):
         return False
 
@@ -194,16 +181,11 @@ class ExecutionList(TopologicalSort):
     def is_cached(self, node_id):
         return self.output_cache.get(node_id) is not None
 
-    async def stage_node_execution(self):
+    def stage_node_execution(self):
         assert self.staged_node_id is None
         if self.is_empty():
             return None, None, None
         available = self.get_ready_nodes()
-        while len(available) == 0 and self.externalBlocks > 0:
-            # Wait for an external block to be released
-            await self.unblockedEvent.wait()
-            self.unblockedEvent.clear()
-            available = self.get_ready_nodes()
         if len(available) == 0:
             cycled_nodes = self.get_nodes_in_cycle()
             # Because cycles composed entirely of static nodes are caught during initial validation,

comfy_execution/progress.py

@@ -1,288 +0,0 @@
-from typing import TypedDict, Dict, Optional
-from typing_extensions import override
-from PIL import Image
-from enum import Enum
-from abc import ABC
-from tqdm import tqdm
-from comfy_execution.graph import DynamicPrompt
-from protocol import BinaryEventTypes
-
-class NodeState(Enum):
-    Pending = "pending"
-    Running = "running"
-    Finished = "finished"
-    Error = "error"
-
-class NodeProgressState(TypedDict):
-    """
-    A class to represent the state of a node's progress.
-    """
-    state: NodeState
-    value: float
-    max: float
-
-class ProgressHandler(ABC):
-    """
-    Abstract base class for progress handlers.
-    Progress handlers receive progress updates and display them in various ways.
-    """
-    def __init__(self, name: str):
-        self.name = name
-        self.enabled = True
-
-    def set_registry(self, registry: "ProgressRegistry"):
-        pass
-
-    def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        """Called when a node starts processing"""
-        pass
-
-    def update_handler(self, node_id: str, value: float, max_value: float,
-                      state: NodeProgressState, prompt_id: str, image: Optional[Image.Image] = None):
-        """Called when a node's progress is updated"""
-        pass
-
-    def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        """Called when a node finishes processing"""
-        pass
-
-    def reset(self):
-        """Called when the progress registry is reset"""
-        pass
-
-    def enable(self):
-        """Enable this handler"""
-        self.enabled = True
-
-    def disable(self):
-        """Disable this handler"""
-        self.enabled = False
-
-class CLIProgressHandler(ProgressHandler):
-    """
-    Handler that displays progress using tqdm progress bars in the CLI.
-    """
-    def __init__(self):
-        super().__init__("cli")
-        self.progress_bars: Dict[str, tqdm] = {}
-
-    @override
-    def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Create a new tqdm progress bar
-        if node_id not in self.progress_bars:
-            self.progress_bars[node_id] = tqdm(
-                total=state["max"],
-                desc=f"Node {node_id}",
-                unit="steps",
-                leave=True,
-                position=len(self.progress_bars)
-            )
-
-    @override
-    def update_handler(self, node_id: str, value: float, max_value: float,
-                      state: NodeProgressState, prompt_id: str, image: Optional[Image.Image] = None):
-        # Handle case where start_handler wasn't called
-        if node_id not in self.progress_bars:
-            self.progress_bars[node_id] = tqdm(
-                total=max_value,
-                desc=f"Node {node_id}",
-                unit="steps",
-                leave=True,
-                position=len(self.progress_bars)
-            )
-            self.progress_bars[node_id].update(value)
-        else:
-            # Update existing progress bar
-            if max_value != self.progress_bars[node_id].total:
-                self.progress_bars[node_id].total = max_value
-            # Calculate the update amount (difference from current position)
-            current_position = self.progress_bars[node_id].n
-            update_amount = value - current_position
-            if update_amount > 0:
-                self.progress_bars[node_id].update(update_amount)
-
-    @override
-    def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Complete and close the progress bar if it exists
-        if node_id in self.progress_bars:
-            # Ensure the bar shows 100% completion
-            remaining = state["max"] - self.progress_bars[node_id].n
-            if remaining > 0:
-                self.progress_bars[node_id].update(remaining)
-            self.progress_bars[node_id].close()
-            del self.progress_bars[node_id]
-
-    @override
-    def reset(self):
-        # Close all progress bars
-        for bar in self.progress_bars.values():
-            bar.close()
-        self.progress_bars.clear()
-
-class WebUIProgressHandler(ProgressHandler):
-    """
-    Handler that sends progress updates to the WebUI via WebSockets.
-    """
-    def __init__(self, server_instance):
-        super().__init__("webui")
-        self.server_instance = server_instance
-
-    def set_registry(self, registry: "ProgressRegistry"):
-        self.registry = registry
-
-    def _send_progress_state(self, prompt_id: str, nodes: Dict[str, NodeProgressState]):
-        """Send the current progress state to the client"""
-        if self.server_instance is None:
-            return
-
-        # Only send info for non-pending nodes
-        active_nodes = {
-            node_id: {
-                "value": state["value"],
-                "max": state["max"],
-                "state": state["state"].value,
-                "node_id": node_id,
-                "prompt_id": prompt_id,
-                "display_node_id": self.registry.dynprompt.get_display_node_id(node_id),
-                "parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
-                "real_node_id": self.registry.dynprompt.get_real_node_id(node_id)
-            }
-            for node_id, state in nodes.items()
-            if state["state"] != NodeState.Pending
-        }
-
-        # Send a combined progress_state message with all node states
-        self.server_instance.send_sync("progress_state", {
-            "prompt_id": prompt_id,
-            "nodes": active_nodes
-        })
-
-    @override
-    def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Send progress state of all nodes
-        if self.registry:
-            self._send_progress_state(prompt_id, self.registry.nodes)
-
-    @override
-    def update_handler(self, node_id: str, value: float, max_value: float,
-                      state: NodeProgressState, prompt_id: str, image: Optional[Image.Image] = None):
-        # Send progress state of all nodes
-        if self.registry:
-            self._send_progress_state(prompt_id, self.registry.nodes)
-        if image:
-            metadata = {
-                "node_id": node_id,
-                "prompt_id": prompt_id,
-                "display_node_id": self.registry.dynprompt.get_display_node_id(node_id),
-                "parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
-                "real_node_id": self.registry.dynprompt.get_real_node_id(node_id)
-            }
-            self.server_instance.send_sync(BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA, (image, metadata), self.server_instance.client_id)
-
-
-    @override
-    def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Send progress state of all nodes
-        if self.registry:
-            self._send_progress_state(prompt_id, self.registry.nodes)
-
-class ProgressRegistry:
-    """
-    Registry that maintains node progress state and notifies registered handlers.
-    """
-    def __init__(self, prompt_id: str, dynprompt: DynamicPrompt):
-        self.prompt_id = prompt_id
-        self.dynprompt = dynprompt
-        self.nodes: Dict[str, NodeProgressState] = {}
-        self.handlers: Dict[str, ProgressHandler] = {}
-
-    def register_handler(self, handler: ProgressHandler) -> None:
-        """Register a progress handler"""
-        self.handlers[handler.name] = handler
-
-    def unregister_handler(self, handler_name: str) -> None:
-        """Unregister a progress handler"""
-        if handler_name in self.handlers:
-            # Allow handler to clean up resources
-            self.handlers[handler_name].reset()
-            del self.handlers[handler_name]
-
-    def enable_handler(self, handler_name: str) -> None:
-        """Enable a progress handler"""
-        if handler_name in self.handlers:
-            self.handlers[handler_name].enable()
-
-    def disable_handler(self, handler_name: str) -> None:
-        """Disable a progress handler"""
-        if handler_name in self.handlers:
-            self.handlers[handler_name].disable()
-
-    def ensure_entry(self, node_id: str) -> NodeProgressState:
-        """Ensure a node entry exists"""
-        if node_id not in self.nodes:
-            self.nodes[node_id] = NodeProgressState(
-                state = NodeState.Pending,
-                value = 0,
-                max = 1
-            )
-        return self.nodes[node_id]
-
-    def start_progress(self, node_id: str) -> None:
-        """Start progress tracking for a node"""
-        entry = self.ensure_entry(node_id)
-        entry["state"] = NodeState.Running
-        entry["value"] = 0.0
-        entry["max"] = 1.0
-
-        # Notify all enabled handlers
-        for handler in self.handlers.values():
-            if handler.enabled:
-                handler.start_handler(node_id, entry, self.prompt_id)
-
-    def update_progress(self, node_id: str, value: float, max_value: float, image: Optional[Image.Image]) -> None:
-        """Update progress for a node"""
-        entry = self.ensure_entry(node_id)
-        entry["state"] = NodeState.Running
-        entry["value"] = value
-        entry["max"] = max_value
-
-        # Notify all enabled handlers
-        for handler in self.handlers.values():
-            if handler.enabled:
-                handler.update_handler(node_id, value, max_value, entry, self.prompt_id, image)
-
-    def finish_progress(self, node_id: str) -> None:
-        """Finish progress tracking for a node"""
-        entry = self.ensure_entry(node_id)
-        entry["state"] = NodeState.Finished
-        entry["value"] = entry["max"]
-
-        # Notify all enabled handlers
-        for handler in self.handlers.values():
-            if handler.enabled:
-                handler.finish_handler(node_id, entry, self.prompt_id)
-
-    def reset_handlers(self) -> None:
-        """Reset all handlers"""
-        for handler in self.handlers.values():
-            handler.reset()
-
-# Global registry instance
-global_progress_registry: ProgressRegistry = ProgressRegistry(prompt_id="", dynprompt=DynamicPrompt({}))
-
-def reset_progress_state(prompt_id: str, dynprompt: DynamicPrompt) -> None:
-    global global_progress_registry
-
-    # Reset existing handlers if registry exists
-    if global_progress_registry is not None:
-        global_progress_registry.reset_handlers()
-
-    # Create new registry
-    global_progress_registry = ProgressRegistry(prompt_id, dynprompt)
-
-def add_progress_handler(handler: ProgressHandler) -> None:
-    handler.set_registry(global_progress_registry)
-    global_progress_registry.register_handler(handler)
-
-def get_progress_state() -> ProgressRegistry:
-    return global_progress_registry

comfy_execution/utils.py

@@ -1,46 +0,0 @@
-import contextvars
-from typing import Optional, NamedTuple
-
-class ExecutionContext(NamedTuple):
-    """
-    Context information about the currently executing node.
-
-    Attributes:
-        node_id: The ID of the currently executing node
-        list_index: The index in a list being processed (for operations on batches/lists)
-    """
-    prompt_id: str
-    node_id: str
-    list_index: Optional[int]
-
-current_executing_context: contextvars.ContextVar[Optional[ExecutionContext]] = contextvars.ContextVar("current_executing_context", default=None)
-
-def get_executing_context() -> Optional[ExecutionContext]:
-    return current_executing_context.get(None)
-
-class CurrentNodeContext:
-    """
-    Context manager for setting the current executing node context.
-
-    Sets the current_executing_context on enter and resets it on exit.
-
-    Example:
-        with CurrentNodeContext(node_id="123", list_index=0):
-            # Code that should run with the current node context set
-            process_image()
-    """
-    def __init__(self, prompt_id: str, node_id: str, list_index: Optional[int] = None):
-        self.context = ExecutionContext(
-            prompt_id= prompt_id,
-            node_id= node_id,
-            list_index= list_index
-        )
-        self.token = None
-
-    def __enter__(self):
-        self.token = current_executing_context.set(self.context)
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.token is not None:
-            current_executing_context.reset(self.token)

comfy_extras/nodes_cosmos.py

@@ -2,6 +2,7 @@ import nodes
 import torch
 import comfy.model_management
 import comfy.utils
+import comfy.latent_formats
 
 
 class EmptyCosmosLatentVideo:
@@ -75,8 +76,53 @@ class CosmosImageToVideoLatent:
         out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1))
         return (out_latent,)
 
+class CosmosPredict2ImageToVideoLatent:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"vae": ("VAE", ),
+                             "width": ("INT", {"default": 848, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+                             "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+                             "length": ("INT", {"default": 93, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}),
+                             "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
+                },
+                "optional": {"start_image": ("IMAGE", ),
+                             "end_image": ("IMAGE", ),
+                }}
+
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "encode"
+
+    CATEGORY = "conditioning/inpaint"
+
+    def encode(self, vae, width, height, length, batch_size, start_image=None, end_image=None):
+        latent = torch.zeros([1, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
+        if start_image is None and end_image is None:
+            out_latent = {}
+            out_latent["samples"] = latent
+            return (out_latent,)
+
+        mask = torch.ones([latent.shape[0], 1, ((length - 1) // 4) + 1, latent.shape[-2], latent.shape[-1]], device=comfy.model_management.intermediate_device())
+
+        if start_image is not None:
+            latent_temp = vae_encode_with_padding(vae, start_image, width, height, length, padding=1)
+            latent[:, :, :latent_temp.shape[-3]] = latent_temp
+            mask[:, :, :latent_temp.shape[-3]] *= 0.0
+
+        if end_image is not None:
+            latent_temp = vae_encode_with_padding(vae, end_image, width, height, length, padding=0)
+            latent[:, :, -latent_temp.shape[-3]:] = latent_temp
+            mask[:, :, -latent_temp.shape[-3]:] *= 0.0
+
+        out_latent = {}
+        latent_format = comfy.latent_formats.Wan21()
+        latent = latent_format.process_out(latent) * mask + latent * (1.0 - mask)
+        out_latent["samples"] = latent.repeat((batch_size, ) + (1,) * (latent.ndim - 1))
+        out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1))
+        return (out_latent,)
 
 NODE_CLASS_MAPPINGS = {
     "EmptyCosmosLatentVideo": EmptyCosmosLatentVideo,
     "CosmosImageToVideoLatent": CosmosImageToVideoLatent,
+    "CosmosPredict2ImageToVideoLatent": CosmosPredict2ImageToVideoLatent,
 }

comfy_extras/nodes_train.py

@@ -552,6 +552,9 @@ class LoraModelLoader:
 
 
 class SaveLoRA:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+
     @classmethod
     def INPUT_TYPES(s):
         return {
@@ -565,7 +568,7 @@ class SaveLoRA:
                 "prefix": (
                     "STRING",
                     {
-                        "default": "trained_lora",
+                        "default": "loras/ComfyUI_trained_lora",
                         "tooltip": "The prefix to use for the saved LoRA file.",
                     },
                 ),
@@ -588,12 +591,13 @@ class SaveLoRA:
     OUTPUT_NODE = True
 
     def save(self, lora, prefix, steps=None):
-        date = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(prefix, self.output_dir)
         if steps is None:
-            output_file = f"models/loras/{prefix}_{date}_lora.safetensors"
+            output_checkpoint = f"{filename}_{counter:05}_.safetensors"
         else:
-            output_file = f"models/loras/{prefix}_{steps}_steps_{date}_lora.safetensors"
-        safetensors.torch.save_file(lora, output_file)
+            output_checkpoint = f"{filename}_{steps}_steps_{counter:05}_.safetensors"
+        output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
+        safetensors.torch.save_file(lora, output_checkpoint)
         return {}
 
 

comfyui_version.py

@@ -1,3 +1,3 @@
 # This file is automatically generated by the build process when version is
 # updated in pyproject.toml.
-__version__ = "0.3.40"
+__version__ = "0.3.41"

execution.py

@@ -8,14 +8,12 @@ import time
 import traceback
 from enum import Enum
 from typing import List, Literal, NamedTuple, Optional
-import asyncio
 
 import torch
 
 import comfy.model_management
 import nodes
 from comfy_execution.caching import (
-    BasicCache,
     CacheKeySetID,
     CacheKeySetInputSignature,
     DependencyAwareCache,
@@ -30,8 +28,6 @@ from comfy_execution.graph import (
 )
 from comfy_execution.graph_utils import GraphBuilder, is_link
 from comfy_execution.validation import validate_node_input
-from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler
-from comfy_execution.utils import CurrentNodeContext
 
 
 class ExecutionResult(Enum):
@@ -43,13 +39,12 @@ class DuplicateNodeError(Exception):
     pass
 
 class IsChangedCache:
-    def __init__(self, prompt_id: str, dynprompt: DynamicPrompt, outputs_cache: BasicCache):
-        self.prompt_id = prompt_id
+    def __init__(self, dynprompt, outputs_cache):
         self.dynprompt = dynprompt
         self.outputs_cache = outputs_cache
         self.is_changed = {}
 
-    async def get(self, node_id):
+    def get(self, node_id):
         if node_id in self.is_changed:
             return self.is_changed[node_id]
 
@@ -67,8 +62,7 @@ class IsChangedCache:
         # Intentionally do not use cached outputs here. We only want constants in IS_CHANGED
         input_data_all, _ = get_input_data(node["inputs"], class_def, node_id, None)
         try:
-            is_changed = await _async_map_node_over_list(self.prompt_id, node_id, class_def, input_data_all, "IS_CHANGED")
-            is_changed = await resolve_map_node_over_list_results(is_changed)
+            is_changed = _map_node_over_list(class_def, input_data_all, "IS_CHANGED")
             node["is_changed"] = [None if isinstance(x, ExecutionBlocker) else x for x in is_changed]
         except Exception as e:
             logging.warning("WARNING: {}".format(e))
@@ -170,19 +164,7 @@ def get_input_data(inputs, class_def, unique_id, outputs=None, dynprompt=None, e
 
 map_node_over_list = None #Don't hook this please
 
-async def resolve_map_node_over_list_results(results):
-    remaining = [x for x in results if isinstance(x, asyncio.Task) and not x.done()]
-    if len(remaining) == 0:
-        return [x.result() if isinstance(x, asyncio.Task) else x for x in results]
-    else:
-        done, pending = await asyncio.wait(remaining)
-        for task in done:
-            exc = task.exception()
-            if exc is not None:
-                raise exc
-        return [x.result() if isinstance(x, asyncio.Task) else x for x in results]
-
-async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None):
+def _map_node_over_list(obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None):
     # check if node wants the lists
     input_is_list = getattr(obj, "INPUT_IS_LIST", False)
 
@@ -196,7 +178,7 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
         return {k: v[i if len(v) > i else -1] for k, v in d.items()}
 
     results = []
-    async def process_inputs(inputs, index=None, input_is_list=False):
+    def process_inputs(inputs, index=None, input_is_list=False):
         if allow_interrupt:
             nodes.before_node_execution()
         execution_block = None
@@ -212,37 +194,20 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
         if execution_block is None:
             if pre_execute_cb is not None and index is not None:
                 pre_execute_cb(index)
-            f = getattr(obj, func)
-            if inspect.iscoroutinefunction(f):
-                async def async_wrapper(f, prompt_id, unique_id, list_index, args):
-                    with CurrentNodeContext(prompt_id, unique_id, list_index):
-                        return await f(**args)
-                task = asyncio.create_task(async_wrapper(f, prompt_id, unique_id, index, args=inputs))
-                # Give the task a chance to execute without yielding
-                await asyncio.sleep(0)
-                if task.done():
-                    result = task.result()
-                    results.append(result)
-                else:
-                    results.append(task)
-            else:
-                with CurrentNodeContext(prompt_id, unique_id, index):
-                    result = f(**inputs)
-                results.append(result)
+            results.append(getattr(obj, func)(**inputs))
         else:
             results.append(execution_block)
 
     if input_is_list:
-        await process_inputs(input_data_all, 0, input_is_list=input_is_list)
+        process_inputs(input_data_all, 0, input_is_list=input_is_list)
     elif max_len_input == 0:
-        await process_inputs({})
+        process_inputs({})
     else:
         for i in range(max_len_input):
             input_dict = slice_dict(input_data_all, i)
-            await process_inputs(input_dict, i)
+            process_inputs(input_dict, i)
     return results
 
-
 def merge_result_data(results, obj):
     # check which outputs need concatenating
     output = []
@@ -264,18 +229,11 @@ def merge_result_data(results, obj):
             output.append([o[i] for o in results])
     return output
 
-async def get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=None, pre_execute_cb=None):
-    return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
-    has_pending_task = any(isinstance(r, asyncio.Task) and not r.done() for r in return_values)
-    if has_pending_task:
-        return return_values, {}, False, has_pending_task
-    output, ui, has_subgraph = get_output_from_returns(return_values, obj)
-    return output, ui, has_subgraph, False
-
-def get_output_from_returns(return_values, obj):
+def get_output_data(obj, input_data_all, execution_block_cb=None, pre_execute_cb=None):
     results = []
     uis = []
     subgraph_results = []
+    return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
     has_subgraph = False
     for i in range(len(return_values)):
         r = return_values[i]
@@ -309,10 +267,6 @@ def get_output_from_returns(return_values, obj):
     else:
         output = []
     ui = dict()
-    # TODO: Think there's an existing bug here
-    # If we're performing a subgraph expansion, we probably shouldn't be returning UI values yet.
-    # They'll get cached without the completed subgraphs. It's an edge case and I'm not aware of
-    # any nodes that use both subgraph expansion and custom UI outputs, but might be a problem in the future.
     if len(uis) > 0:
         ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()}
     return output, ui, has_subgraph
@@ -325,7 +279,7 @@ def format_value(x):
     else:
         return str(x)
 
-async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes):
+def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results):
     unique_id = current_item
     real_node_id = dynprompt.get_real_node_id(unique_id)
     display_node_id = dynprompt.get_display_node_id(unique_id)
@@ -337,16 +291,11 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
         if server.client_id is not None:
             cached_output = caches.ui.get(unique_id) or {}
             server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": cached_output.get("output",None), "prompt_id": prompt_id }, server.client_id)
-        get_progress_state().finish_progress(unique_id)
         return (ExecutionResult.SUCCESS, None, None)
 
     input_data_all = None
     try:
-        if unique_id in pending_async_nodes:
-            results = [r.result() if isinstance(r, asyncio.Task) else r for r in pending_async_nodes[unique_id]]
-            del pending_async_nodes[unique_id]
-            output_data, output_ui, has_subgraph = get_output_from_returns(results, class_def)
-        elif unique_id in pending_subgraph_results:
+        if unique_id in pending_subgraph_results:
             cached_results = pending_subgraph_results[unique_id]
             resolved_outputs = []
             for is_subgraph, result in cached_results:
@@ -368,7 +317,6 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
             output_ui = []
             has_subgraph = False
         else:
-            get_progress_state().start_progress(unique_id)
             input_data_all, missing_keys = get_input_data(inputs, class_def, unique_id, caches.outputs, dynprompt, extra_data)
             if server.client_id is not None:
                 server.last_node_id = display_node_id
@@ -380,8 +328,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                 caches.objects.set(unique_id, obj)
 
             if hasattr(obj, "check_lazy_status"):
-                required_inputs = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, "check_lazy_status", allow_interrupt=True)
-                required_inputs = await resolve_map_node_over_list_results(required_inputs)
+                required_inputs = _map_node_over_list(obj, input_data_all, "check_lazy_status", allow_interrupt=True)
                 required_inputs = set(sum([r for r in required_inputs if isinstance(r,list)], []))
                 required_inputs = [x for x in required_inputs if isinstance(x,str) and (
                     x not in input_data_all or x in missing_keys
@@ -410,18 +357,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                 else:
                     return block
             def pre_execute_cb(call_index):
-                # TODO - How to handle this with async functions without contextvars (which requires Python 3.12)?
                 GraphBuilder.set_default_prefix(unique_id, call_index, 0)
-            output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
-            if has_pending_tasks:
-                pending_async_nodes[unique_id] = output_data
-                unblock = execution_list.add_external_block(unique_id)
-                async def await_completion():
-                    tasks = [x for x in output_data if isinstance(x, asyncio.Task)]
-                    await asyncio.gather(*tasks)
-                    unblock()
-                asyncio.create_task(await_completion())
-                return (ExecutionResult.PENDING, None, None)
+            output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
         if len(output_ui) > 0:
             caches.ui.set(unique_id, {
                 "meta": {
@@ -464,8 +401,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                     cached_outputs.append((True, node_outputs))
             new_node_ids = set(new_node_ids)
             for cache in caches.all:
-                subcache = await cache.ensure_subcache_for(unique_id, new_node_ids)
-                subcache.clean_unused()
+                cache.ensure_subcache_for(unique_id, new_node_ids).clean_unused()
             for node_id in new_output_ids:
                 execution_list.add_node(node_id)
             for link in new_output_links:
@@ -507,7 +443,6 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
 
         return (ExecutionResult.FAILURE, error_details, ex)
 
-    get_progress_state().finish_progress(unique_id)
     executed.add(unique_id)
 
     return (ExecutionResult.SUCCESS, None, None)
@@ -562,11 +497,6 @@ class PromptExecutor:
             self.add_message("execution_error", mes, broadcast=False)
 
     def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
-        asyncio_loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(asyncio_loop)
-        asyncio.run(self.execute_async(prompt, prompt_id, extra_data, execute_outputs))
-
-    async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
         nodes.interrupt_processing(False)
 
         if "client_id" in extra_data:
@@ -579,11 +509,9 @@ class PromptExecutor:
 
         with torch.inference_mode():
             dynamic_prompt = DynamicPrompt(prompt)
-            reset_progress_state(prompt_id, dynamic_prompt)
-            add_progress_handler(WebUIProgressHandler(self.server))
-            is_changed_cache = IsChangedCache(prompt_id, dynamic_prompt, self.caches.outputs)
+            is_changed_cache = IsChangedCache(dynamic_prompt, self.caches.outputs)
             for cache in self.caches.all:
-                await cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache)
+                cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache)
                 cache.clean_unused()
 
             cached_nodes = []
@@ -596,7 +524,6 @@ class PromptExecutor:
                           { "nodes": cached_nodes, "prompt_id": prompt_id},
                           broadcast=False)
             pending_subgraph_results = {}
-            pending_async_nodes = {} # TODO - Unify this with pending_subgraph_results
             executed = set()
             execution_list = ExecutionList(dynamic_prompt, self.caches.outputs)
             current_outputs = self.caches.outputs.all_node_ids()
@@ -604,13 +531,12 @@ class PromptExecutor:
                 execution_list.add_node(node_id)
 
             while not execution_list.is_empty():
-                node_id, error, ex = await execution_list.stage_node_execution()
+                node_id, error, ex = execution_list.stage_node_execution()
                 if error is not None:
                     self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
                     break
 
-                assert node_id is not None, "Node ID should not be None at this point"
-                result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes)
+                result, error, ex = execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results)
                 self.success = result != ExecutionResult.FAILURE
                 if result == ExecutionResult.FAILURE:
                     self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
@@ -640,7 +566,7 @@ class PromptExecutor:
                 comfy.model_management.unload_all_models()
 
 
-async def validate_inputs(prompt_id, prompt, item, validated):
+def validate_inputs(prompt, item, validated):
     unique_id = item
     if unique_id in validated:
         return validated[unique_id]
@@ -717,7 +643,7 @@ async def validate_inputs(prompt_id, prompt, item, validated):
                 errors.append(error)
                 continue
             try:
-                r = await validate_inputs(prompt_id, prompt, o_id, validated)
+                r = validate_inputs(prompt, o_id, validated)
                 if r[0] is False:
                     # `r` will be set in `validated[o_id]` already
                     valid = False
@@ -842,8 +768,7 @@ async def validate_inputs(prompt_id, prompt, item, validated):
             input_filtered['input_types'] = [received_types]
 
         #ret = obj_class.VALIDATE_INPUTS(**input_filtered)
-        ret = await _async_map_node_over_list(prompt_id, unique_id, obj_class, input_filtered, "VALIDATE_INPUTS")
-        ret = await resolve_map_node_over_list_results(ret)
+        ret = _map_node_over_list(obj_class, input_filtered, "VALIDATE_INPUTS")
         for x in input_filtered:
             for i, r in enumerate(ret):
                 if r is not True and not isinstance(r, ExecutionBlocker):
@@ -876,7 +801,7 @@ def full_type_name(klass):
         return klass.__qualname__
     return module + '.' + klass.__qualname__
 
-async def validate_prompt(prompt_id, prompt):
+def validate_prompt(prompt):
     outputs = set()
     for x in prompt:
         if 'class_type' not in prompt[x]:
@@ -919,7 +844,7 @@ async def validate_prompt(prompt_id, prompt):
         valid = False
         reasons = []
         try:
-            m = await validate_inputs(prompt_id, prompt, o, validated)
+            m = validate_inputs(prompt, o, validated)
             valid = m[0]
             reasons = m[1]
         except Exception as ex:

main.py

@@ -11,8 +11,6 @@ import itertools
 import utils.extra_config
 import logging
 import sys
-from comfy_execution.progress import get_progress_state
-from comfy_execution.utils import get_executing_context
 
 if __name__ == "__main__":
     #NOTE: These do not do anything on core ComfyUI, they are for custom nodes.
@@ -130,7 +128,7 @@ import comfy.utils
 
 import execution
 import server
-from protocol import BinaryEventTypes
+from server import BinaryEventTypes
 import nodes
 import comfy.model_management
 import comfyui_version
@@ -220,25 +218,14 @@ async def run(server_instance, address='', port=8188, verbose=True, call_on_star
         server_instance.start_multi_address(addresses, call_on_start, verbose), server_instance.publish_loop()
     )
 
+
 def hijack_progress(server_instance):
-    def hook(value, total, preview_image, prompt_id=None, node_id=None):
-        executing_context = get_executing_context()
-        if prompt_id is None and executing_context is not None:
-            prompt_id = executing_context.prompt_id
-        if node_id is None and executing_context is not None:
-            node_id = executing_context.node_id
+    def hook(value, total, preview_image):
         comfy.model_management.throw_exception_if_processing_interrupted()
-        if prompt_id is None:
-            prompt_id = server_instance.last_prompt_id
-        if node_id is None:
-            node_id = server_instance.last_node_id
-        progress = {"value": value, "max": total, "prompt_id": prompt_id, "node": node_id}
-        get_progress_state().update_progress(node_id, value, total, preview_image)
+        progress = {"value": value, "max": total, "prompt_id": server_instance.last_prompt_id, "node": server_instance.last_node_id}
 
         server_instance.send_sync("progress", progress, server_instance.client_id)
         if preview_image is not None:
-            # Also send old method for backward compatibility
-            # TODO - Remove after this repo is updated to frontend with metadata support
             server_instance.send_sync(BinaryEventTypes.UNENCODED_PREVIEW_IMAGE, preview_image, server_instance.client_id)
 
     comfy.utils.set_progress_bar_global_hook(hook)

protocol.py

@@ -1,7 +0,0 @@
-
-class BinaryEventTypes:
-    PREVIEW_IMAGE = 1
-    UNENCODED_PREVIEW_IMAGE = 2
-    TEXT = 3
-    PREVIEW_IMAGE_WITH_METADATA = 4
-

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI"
-version = "0.3.40"
+version = "0.3.41"
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.9"

server.py

@@ -35,7 +35,11 @@ from app.model_manager import ModelFileManager
 from app.custom_node_manager import CustomNodeManager
 from typing import Optional, Union
 from api_server.routes.internal.internal_routes import InternalRoutes
-from protocol import BinaryEventTypes
+
+class BinaryEventTypes:
+    PREVIEW_IMAGE = 1
+    UNENCODED_PREVIEW_IMAGE = 2
+    TEXT = 3
 
 async def send_socket_catch_exception(function, message):
     try:
@@ -639,8 +643,7 @@ class PromptServer():
 
             if "prompt" in json_data:
                 prompt = json_data["prompt"]
-                prompt_id = str(uuid.uuid4())
-                valid = await execution.validate_prompt(prompt_id, prompt)
+                valid = execution.validate_prompt(prompt)
                 extra_data = {}
                 if "extra_data" in json_data:
                     extra_data = json_data["extra_data"]
@@ -648,6 +651,7 @@ class PromptServer():
                 if "client_id" in json_data:
                     extra_data["client_id"] = json_data["client_id"]
                 if valid[0]:
+                    prompt_id = str(uuid.uuid4())
                     outputs_to_execute = valid[2]
                     self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
                     response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
@@ -762,10 +766,6 @@ class PromptServer():
     async def send(self, event, data, sid=None):
         if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE:
             await self.send_image(data, sid=sid)
-        elif event == BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA:
-            # data is (preview_image, metadata)
-            preview_image, metadata = data
-            await self.send_image_with_metadata(preview_image, metadata, sid=sid)
         elif isinstance(data, (bytes, bytearray)):
             await self.send_bytes(event, data, sid)
         else:
@@ -804,43 +804,6 @@ class PromptServer():
         preview_bytes = bytesIO.getvalue()
         await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid)
 
-    async def send_image_with_metadata(self, image_data, metadata=None, sid=None):
-        image_type = image_data[0]
-        image = image_data[1]
-        max_size = image_data[2]
-        if max_size is not None:
-            if hasattr(Image, 'Resampling'):
-                resampling = Image.Resampling.BILINEAR
-            else:
-                resampling = Image.Resampling.LANCZOS
-
-            image = ImageOps.contain(image, (max_size, max_size), resampling)
-
-        mimetype = "image/png" if image_type == "PNG" else "image/jpeg"
-
-        # Prepare metadata
-        if metadata is None:
-            metadata = {}
-        metadata["image_type"] = mimetype
-
-        # Serialize metadata as JSON
-        import json
-        metadata_json = json.dumps(metadata).encode('utf-8')
-        metadata_length = len(metadata_json)
-
-        # Prepare image data
-        bytesIO = BytesIO()
-        image.save(bytesIO, format=image_type, quality=95, compress_level=1)
-        image_bytes = bytesIO.getvalue()
-
-        # Combine metadata and image
-        combined_data = bytearray()
-        combined_data.extend(struct.pack(">I", metadata_length))
-        combined_data.extend(metadata_json)
-        combined_data.extend(image_bytes)
-
-        await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA, combined_data, sid=sid)
-
     async def send_bytes(self, event, data, sid=None):
         message = self.encode_bytes(event, data)
 

tests-unit/requirements.txt

@@ -1,4 +1,3 @@
 pytest>=7.8.0
 pytest-aiohttp
 pytest-asyncio
-websocket-client

tests/inference/extra_model_paths.yaml

@@ -1,4 +1,4 @@
 # Config for testing nodes
 testing:
-    custom_nodes: testing_nodes
+    custom_nodes: tests/inference/testing_nodes
 

tests/inference/test_execution.py

@@ -252,7 +252,7 @@ class TestExecution:
 
     @pytest.mark.parametrize("test_type, test_value", [
         ("StubInt", 5),
-        ("StubMask", 5.0)
+        ("StubFloat", 5.0)
     ])
     def test_validation_error_edge1(self, test_type, test_value, client: ComfyClient, builder: GraphBuilder):
         g = builder
@@ -497,69 +497,6 @@ class TestExecution:
         assert numpy.array(images[0]).min() == 63 and numpy.array(images[0]).max() == 63, "Image should have value 0.25"
         assert not result.did_run(test_node), "The execution should have been cached"
 
-    def test_parallel_sleep_nodes(self, client: ComfyClient, builder: GraphBuilder):
-        g = builder
-        image = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
-
-        # Create sleep nodes for each duration
-        sleep_node1 = g.node("TestSleep", value=image.out(0), seconds=2.8)
-        sleep_node2 = g.node("TestSleep", value=image.out(0), seconds=2.9)
-        sleep_node3 = g.node("TestSleep", value=image.out(0), seconds=3.0)
-
-        # Add outputs to verify the execution
-        _output1 = g.node("PreviewImage", images=sleep_node1.out(0))
-        _output2 = g.node("PreviewImage", images=sleep_node2.out(0))
-        _output3 = g.node("PreviewImage", images=sleep_node3.out(0))
-
-        start_time = time.time()
-        result = client.run(g)
-        elapsed_time = time.time() - start_time
-
-        # The test should take around 0.4 seconds (the longest sleep duration)
-        # plus some overhead, but definitely less than the sum of all sleeps (0.9s)
-        # We'll allow for up to 0.8s total to account for overhead
-        assert elapsed_time < 4.0, f"Parallel execution took {elapsed_time}s, expected less than 0.8s"
-
-        # Verify that all nodes executed
-        assert result.did_run(sleep_node1), "Sleep node 1 should have run"
-        assert result.did_run(sleep_node2), "Sleep node 2 should have run"
-        assert result.did_run(sleep_node3), "Sleep node 3 should have run"
-
-    def test_parallel_sleep_expansion(self, client: ComfyClient, builder: GraphBuilder):
-        g = builder
-        # Create input images with different values
-        image1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
-        image2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
-        image3 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
-
-        # Create a TestParallelSleep node that expands into multiple TestSleep nodes
-        parallel_sleep = g.node("TestParallelSleep",
-                                image1=image1.out(0),
-                                image2=image2.out(0),
-                                image3=image3.out(0),
-                                sleep1=0.4,
-                                sleep2=0.5,
-                                sleep3=0.6)
-        output = g.node("SaveImage", images=parallel_sleep.out(0))
-
-        start_time = time.time()
-        result = client.run(g)
-        elapsed_time = time.time() - start_time
-
-        # Similar to the previous test, expect parallel execution of the sleep nodes
-        # which should complete in less than the sum of all sleeps
-        assert elapsed_time < 0.8, f"Expansion execution took {elapsed_time}s, expected less than 0.8s"
-
-        # Verify the parallel sleep node executed
-        assert result.did_run(parallel_sleep), "ParallelSleep node should have run"
-
-        # Verify we get an image as output (blend of the three input images)
-        result_images = result.get_images(output)
-        assert len(result_images) == 1, "Should have 1 image"
-        # Average pixel value should be around 170 (255 * 2 // 3)
-        avg_value = numpy.array(result_images[0]).mean()
-        assert avg_value == 170, f"Image average value {avg_value} should be 170"
-
     # This tests that nodes with OUTPUT_IS_LIST function correctly when they receive an ExecutionBlocker
     # as input. We also test that when that list (containing an ExecutionBlocker) is passed to a node,
     # only that one entry in the list is blocked.

tests/inference/testing_nodes/testing-pack/specific_tests.py

@@ -1,11 +1,6 @@
 import torch
-import time
-import asyncio
-from comfy.utils import ProgressBar
 from .tools import VariantSupport
 from comfy_execution.graph_utils import GraphBuilder
-from comfy.comfy_types.node_typing import ComfyNodeABC
-from comfy.comfy_types import IO
 
 class TestLazyMixImages:
     @classmethod
@@ -338,131 +333,6 @@ class TestMixedExpansionReturns:
                 "expand": g.finalize(),
             }
 
-class TestSamplingInExpansion:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "clip": ("CLIP",),
-                "vae": ("VAE",),
-                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                "steps": ("INT", {"default": 20, "min": 1, "max": 100}),
-                "cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 30.0}),
-                "prompt": ("STRING", {"multiline": True, "default": "a beautiful landscape with mountains and trees"}),
-                "negative_prompt": ("STRING", {"multiline": True, "default": "blurry, bad quality, worst quality"}),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "sampling_in_expansion"
-
-    CATEGORY = "Testing/Nodes"
-
-    def sampling_in_expansion(self, model, clip, vae, seed, steps, cfg, prompt, negative_prompt):
-        g = GraphBuilder()
-
-        # Create a basic image generation workflow using the input model, clip and vae
-        # 1. Setup text prompts using the provided CLIP model
-        positive_prompt = g.node("CLIPTextEncode",
-                               text=prompt,
-                               clip=clip)
-        negative_prompt = g.node("CLIPTextEncode",
-                                text=negative_prompt,
-                                clip=clip)
-
-        # 2. Create empty latent with specified size
-        empty_latent = g.node("EmptyLatentImage", width=512, height=512, batch_size=1)
-
-        # 3. Setup sampler and generate image latent
-        sampler = g.node("KSampler",
-                        model=model,
-                        positive=positive_prompt.out(0),
-                        negative=negative_prompt.out(0),
-                        latent_image=empty_latent.out(0),
-                        seed=seed,
-                        steps=steps,
-                        cfg=cfg,
-                        sampler_name="euler_ancestral",
-                        scheduler="normal")
-
-        # 4. Decode latent to image using VAE
-        output = g.node("VAEDecode", samples=sampler.out(0), vae=vae)
-
-        return {
-            "result": (output.out(0),),
-            "expand": g.finalize(),
-        }
-
-class TestSleep(ComfyNodeABC):
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "value": (IO.ANY, {}),
-                "seconds": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 9999.0, "step": 0.01, "tooltip": "The amount of seconds to sleep."}),
-            },
-            "hidden": {
-                "unique_id": "UNIQUE_ID",
-            },
-        }
-    RETURN_TYPES = (IO.ANY,)
-    FUNCTION = "sleep"
-
-    CATEGORY = "_for_testing"
-
-    async def sleep(self, value, seconds, unique_id):
-        pbar = ProgressBar(seconds, node_id=unique_id)
-        start = time.time()
-        expiration = start + seconds
-        now = start
-        while now < expiration:
-            now = time.time()
-            pbar.update_absolute(now - start)
-            await asyncio.sleep(0.01)
-        return (value,)
-
-class TestParallelSleep(ComfyNodeABC):
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "image1": ("IMAGE", ),
-                "image2": ("IMAGE", ),
-                "image3": ("IMAGE", ),
-                "sleep1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.01}),
-                "sleep2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.01}),
-                "sleep3": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step": 0.01}),
-            },
-            "hidden": {
-                "unique_id": "UNIQUE_ID",
-            },
-        }
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "parallel_sleep"
-    CATEGORY = "_for_testing"
-    OUTPUT_NODE = True
-
-    def parallel_sleep(self, image1, image2, image3, sleep1, sleep2, sleep3, unique_id):
-        # Create a graph dynamically with three TestSleep nodes
-        g = GraphBuilder()
-
-        # Create sleep nodes for each duration and image
-        sleep_node1 = g.node("TestSleep", value=image1, seconds=sleep1)
-        sleep_node2 = g.node("TestSleep", value=image2, seconds=sleep2)
-        sleep_node3 = g.node("TestSleep", value=image3, seconds=sleep3)
-
-        # Blend the results using TestVariadicAverage
-        blend = g.node("TestVariadicAverage",
-                       input1=sleep_node1.out(0),
-                       input2=sleep_node2.out(0),
-                       input3=sleep_node3.out(0))
-
-        return {
-            "result": (blend.out(0),),
-            "expand": g.finalize(),
-        }
-
 TEST_NODE_CLASS_MAPPINGS = {
     "TestLazyMixImages": TestLazyMixImages,
     "TestVariadicAverage": TestVariadicAverage,
@@ -475,9 +345,6 @@ TEST_NODE_CLASS_MAPPINGS = {
     "TestCustomValidation5": TestCustomValidation5,
     "TestDynamicDependencyCycle": TestDynamicDependencyCycle,
     "TestMixedExpansionReturns": TestMixedExpansionReturns,
-    "TestSamplingInExpansion": TestSamplingInExpansion,
-    "TestSleep": TestSleep,
-    "TestParallelSleep": TestParallelSleep,
 }
 
 TEST_NODE_DISPLAY_NAME_MAPPINGS = {
@@ -492,7 +359,4 @@ TEST_NODE_DISPLAY_NAME_MAPPINGS = {
     "TestCustomValidation5": "Custom Validation 5",
     "TestDynamicDependencyCycle": "Dynamic Dependency Cycle",
     "TestMixedExpansionReturns": "Mixed Expansion Returns",
-    "TestSamplingInExpansion": "Sampling In Expansion",
-    "TestSleep": "Test Sleep",
-    "TestParallelSleep": "Test Parallel Sleep",
 }