Merge 5884783f90 into 6592bffc60

* Add phi_2 solver type to seeds_2

* Add sampler node of seeds_2

comfy/k_diffusion/sampling.py

@@ -1557,10 +1557,13 @@ def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None
 
 
 @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):
+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, solver_type="phi_1"):
     """SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 2.
     arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
     """
+    if solver_type not in {"phi_1", "phi_2"}:
+        raise ValueError("solver_type must be 'phi_1' or 'phi_2'")
+
     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
@@ -1600,8 +1603,14 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
         denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
 
         # Step 2
-        denoised_d = torch.lerp(denoised, denoised_2, fac)
-        x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * ei_h_phi_1(-h_eta) * denoised_d
+        if solver_type == "phi_1":
+            denoised_d = torch.lerp(denoised, denoised_2, fac)
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * ei_h_phi_1(-h_eta) * denoised_d
+        elif solver_type == "phi_2":
+            b2 = ei_h_phi_2(-h_eta) / r
+            b1 = ei_h_phi_1(-h_eta) - b2
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * (b1 * denoised + b2 * denoised_2)
+
         if inject_noise:
             segment_factor = (r - 1) * h * eta
             sde_noise = sde_noise * segment_factor.exp()

comfy/ldm/chroma/model.py

@@ -183,7 +183,7 @@ class Chroma(nn.Module):
         transformer_options["total_blocks"] = len(self.double_blocks)
         transformer_options["block_type"] = "double"
         for i, block in enumerate(self.double_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if i not in self.skip_mmdit:
                 double_mod = (
                     self.get_modulations(mod_vectors, "double_img", idx=i),
@@ -229,7 +229,7 @@ class Chroma(nn.Module):
         transformer_options["total_blocks"] = len(self.single_blocks)
         transformer_options["block_type"] = "single"
         for i, block in enumerate(self.single_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if i not in self.skip_dit:
                 single_mod = self.get_modulations(mod_vectors, "single", idx=i)
                 if ("single_block", i) in blocks_replace:

comfy/ldm/flux/model.py

@@ -186,7 +186,7 @@ class Flux(nn.Module):
         transformer_options["total_blocks"] = len(self.double_blocks)
         transformer_options["block_type"] = "double"
         for i, block in enumerate(self.double_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -233,7 +233,7 @@ class Flux(nn.Module):
         transformer_options["total_blocks"] = len(self.single_blocks)
         transformer_options["block_type"] = "single"
         for i, block in enumerate(self.single_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if ("single_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}

comfy/ldm/hunyuan_video/model.py

@@ -393,7 +393,7 @@ class HunyuanVideo(nn.Module):
         transformer_options["total_blocks"] = len(self.double_blocks)
         transformer_options["block_type"] = "double"
         for i, block in enumerate(self.double_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -418,7 +418,7 @@ class HunyuanVideo(nn.Module):
         transformer_options["total_blocks"] = len(self.single_blocks)
         transformer_options["block_type"] = "single"
         for i, block in enumerate(self.single_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if ("single_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}

comfy/ldm/kandinsky5/model.py

@@ -375,7 +375,7 @@ class Kandinsky5(nn.Module):
         transformer_options["total_blocks"] = len(self.visual_transformer_blocks)
         transformer_options["block_type"] = "double"
         for i, block in enumerate(self.visual_transformer_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     return block(x=args["x"], context=args["context"], time_embed=args["time_embed"], freqs=args["freqs"], transformer_options=args.get("transformer_options"))

comfy/ldm/modules/attention.py

@@ -921,7 +921,7 @@ class SpatialTransformer(nn.Module):
         if self.use_linear:
             x = self.proj_in(x)
         for i, block in enumerate(self.transformer_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             x = block(x, context=context[i], transformer_options=transformer_options)
         if self.use_linear:
             x = self.proj_out(x)
@@ -1067,7 +1067,7 @@ class SpatialVideoTransformer(SpatialTransformer):
         for it_, (block, mix_block) in enumerate(
             zip(self.transformer_blocks, self.time_stack)
         ):
-            transformer_options["block_index"] = it_
+            transformer_options["block_index"] = torch.tensor(it_, dtype=torch.uint8, device="cpu")
             x = block(
                 x,
                 context=spatial_context,

comfy/ldm/qwen_image/model.py

@@ -442,7 +442,7 @@ class QwenImageTransformer2DModel(nn.Module):
         transformer_options["total_blocks"] = len(self.transformer_blocks)
         transformer_options["block_type"] = "double"
         for i, block in enumerate(self.transformer_blocks):
-            transformer_options["block_index"] = i
+            transformer_options["block_index"] = torch.tensor(i, dtype=torch.uint8, device="cpu")
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}

comfy/ops.py

@@ -592,7 +592,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                     quant_conf = {"format": self.quant_format}
                     if self._full_precision_mm:
                         quant_conf["full_precision_matrix_mult"] = True
-                    sd["{}comfy_quant".format(prefix)] = torch.frombuffer(json.dumps(quant_conf).encode('utf-8'), dtype=torch.uint8)
+                    sd["{}comfy_quant".format(prefix)] = torch.tensor(list(json.dumps(quant_conf).encode('utf-8')), dtype=torch.uint8)
                 return sd
 
             def _forward(self, input, weight, bias):

comfy/utils.py

@@ -1262,6 +1262,6 @@ def convert_old_quants(state_dict, model_prefix="", metadata={}):
     if quant_metadata is not None:
         layers = quant_metadata["layers"]
         for k, v in layers.items():
-            state_dict["{}.comfy_quant".format(k)] = torch.frombuffer(json.dumps(v).encode('utf-8'), dtype=torch.uint8)
+            state_dict["{}.comfy_quant".format(k)] = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8)
 
     return state_dict, metadata

comfy_extras/nodes_custom_sampler.py

@@ -659,6 +659,31 @@ class SamplerSASolver(io.ComfyNode):
     get_sampler = execute
 
 
+class SamplerSEEDS2(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SamplerSEEDS2",
+            category="sampling/custom_sampling/samplers",
+            inputs=[
+                io.Combo.Input("solver_type", options=["phi_1", "phi_2"]),
+                io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, tooltip="Stochastic strength"),
+                io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, tooltip="SDE noise multiplier"),
+                io.Float.Input("r", default=0.5, min=0.01, max=1.0, step=0.01, round=False, tooltip="Relative step size for the intermediate stage (c2 node)"),
+            ],
+            outputs=[io.Sampler.Output()]
+        )
+
+    @classmethod
+    def execute(cls, solver_type, eta, s_noise, r) -> io.NodeOutput:
+        sampler_name = "seeds_2"
+        sampler = comfy.samplers.ksampler(
+            sampler_name,
+            {"eta": eta, "s_noise": s_noise, "r": r, "solver_type": solver_type},
+        )
+        return io.NodeOutput(sampler)
+
+
 class Noise_EmptyNoise:
     def __init__(self):
         self.seed = 0
@@ -996,6 +1021,7 @@ class CustomSamplersExtension(ComfyExtension):
             SamplerDPMAdaptative,
             SamplerER_SDE,
             SamplerSASolver,
+            SamplerSEEDS2,
             SplitSigmas,
             SplitSigmasDenoise,
             FlipSigmas,