Merge 03be9b6efa 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/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_api/feature_flags.py

@@ -5,12 +5,12 @@ This module handles capability negotiation between frontend and backend,
 allowing graceful protocol evolution while maintaining backward compatibility.
 """
 
-from typing import Any, Dict
+from typing import Any
 
 from comfy.cli_args import args
 
 # Default server capabilities
-SERVER_FEATURE_FLAGS: Dict[str, Any] = {
+SERVER_FEATURE_FLAGS: dict[str, Any] = {
     "supports_preview_metadata": True,
     "max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes
     "extension": {"manager": {"supports_v4": True}},
@@ -18,7 +18,7 @@ SERVER_FEATURE_FLAGS: Dict[str, Any] = {
 
 
 def get_connection_feature(
-    sockets_metadata: Dict[str, Dict[str, Any]],
+    sockets_metadata: dict[str, dict[str, Any]],
     sid: str,
     feature_name: str,
     default: Any = False
@@ -42,7 +42,7 @@ def get_connection_feature(
 
 
 def supports_feature(
-    sockets_metadata: Dict[str, Dict[str, Any]],
+    sockets_metadata: dict[str, dict[str, Any]],
     sid: str,
     feature_name: str
 ) -> bool:
@@ -60,7 +60,7 @@ def supports_feature(
     return get_connection_feature(sockets_metadata, sid, feature_name, False) is True
 
 
-def get_server_features() -> Dict[str, Any]:
+def get_server_features() -> dict[str, Any]:
     """
     Get the server's feature flags.
 

comfy_api/internal/api_registry.py

@@ -1,4 +1,4 @@
-from typing import Type, List, NamedTuple
+from typing import NamedTuple
 from comfy_api.internal.singleton import ProxiedSingleton
 from packaging import version as packaging_version
 
@@ -10,7 +10,7 @@ class ComfyAPIBase(ProxiedSingleton):
 
 class ComfyAPIWithVersion(NamedTuple):
     version: str
-    api_class: Type[ComfyAPIBase]
+    api_class: type[ComfyAPIBase]
 
 
 def parse_version(version_str: str) -> packaging_version.Version:
@@ -23,16 +23,16 @@ def parse_version(version_str: str) -> packaging_version.Version:
     return packaging_version.parse(version_str)
 
 
-registered_versions: List[ComfyAPIWithVersion] = []
+registered_versions: list[ComfyAPIWithVersion] = []
 
 
-def register_versions(versions: List[ComfyAPIWithVersion]):
+def register_versions(versions: list[ComfyAPIWithVersion]):
     versions.sort(key=lambda x: parse_version(x.version))
     global registered_versions
     registered_versions = versions
 
 
-def get_all_versions() -> List[ComfyAPIWithVersion]:
+def get_all_versions() -> list[ComfyAPIWithVersion]:
     """
     Returns a list of all registered ComfyAPI versions.
     """

comfy_api/internal/async_to_sync.py

@@ -8,7 +8,7 @@ import os
 import textwrap
 import threading
 from enum import Enum
-from typing import Optional, Type, get_origin, get_args, get_type_hints
+from typing import Optional, get_origin, get_args, get_type_hints
 
 
 class TypeTracker:
@@ -193,7 +193,7 @@ class AsyncToSyncConverter:
         return result_container["result"]
 
     @classmethod
-    def create_sync_class(cls, async_class: Type, thread_pool_size=10) -> Type:
+    def create_sync_class(cls, async_class: type, thread_pool_size=10) -> type:
         """
         Creates a new class with synchronous versions of all async methods.
 
@@ -563,7 +563,7 @@ class AsyncToSyncConverter:
 
     @classmethod
     def _generate_imports(
-        cls, async_class: Type, type_tracker: TypeTracker
+        cls, async_class: type, type_tracker: TypeTracker
     ) -> list[str]:
         """Generate import statements for the stub file."""
         imports = []
@@ -628,7 +628,7 @@ class AsyncToSyncConverter:
         return imports
 
     @classmethod
-    def _get_class_attributes(cls, async_class: Type) -> list[tuple[str, Type]]:
+    def _get_class_attributes(cls, async_class: type) -> list[tuple[str, type]]:
         """Extract class attributes that are classes themselves."""
         class_attributes = []
 
@@ -654,7 +654,7 @@ class AsyncToSyncConverter:
     def _generate_inner_class_stub(
         cls,
         name: str,
-        attr: Type,
+        attr: type,
         indent: str = "    ",
         type_tracker: Optional[TypeTracker] = None,
     ) -> list[str]:
@@ -782,7 +782,7 @@ class AsyncToSyncConverter:
         return processed
 
     @classmethod
-    def generate_stub_file(cls, async_class: Type, sync_class: Type) -> None:
+    def generate_stub_file(cls, async_class: type, sync_class: type) -> None:
         """
         Generate a .pyi stub file for the sync class to help IDEs with type checking.
         """
@@ -988,7 +988,7 @@ class AsyncToSyncConverter:
             logging.error(traceback.format_exc())
 
 
-def create_sync_class(async_class: Type, thread_pool_size=10) -> Type:
+def create_sync_class(async_class: type, thread_pool_size=10) -> type:
     """
     Creates a sync version of an async class
 

comfy_api/internal/singleton.py

@@ -1,4 +1,4 @@
-from typing import Type, TypeVar
+from typing import TypeVar
 
 class SingletonMetaclass(type):
     T = TypeVar("T", bound="SingletonMetaclass")
@@ -11,13 +11,13 @@ class SingletonMetaclass(type):
             )
         return cls._instances[cls]
 
-    def inject_instance(cls: Type[T], instance: T) -> None:
+    def inject_instance(cls: type[T], instance: T) -> None:
         assert cls not in SingletonMetaclass._instances, (
             "Cannot inject instance after first instantiation"
         )
         SingletonMetaclass._instances[cls] = instance
 
-    def get_instance(cls: Type[T], *args, **kwargs) -> T:
+    def get_instance(cls: type[T], *args, **kwargs) -> T:
         """
         Gets the singleton instance of the class, creating it if it doesn't exist.
         """

comfy_api/latest/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Type, TYPE_CHECKING
+from typing import TYPE_CHECKING
 from comfy_api.internal import ComfyAPIBase
 from comfy_api.internal.singleton import ProxiedSingleton
 from comfy_api.internal.async_to_sync import create_sync_class
@@ -113,7 +113,7 @@ ComfyAPI = ComfyAPI_latest
 if TYPE_CHECKING:
     import comfy_api.latest.generated.ComfyAPISyncStub  # type: ignore
 
-    ComfyAPISync: Type[comfy_api.latest.generated.ComfyAPISyncStub.ComfyAPISyncStub]
+    ComfyAPISync: type[comfy_api.latest.generated.ComfyAPISyncStub.ComfyAPISyncStub]
 ComfyAPISync = create_sync_class(ComfyAPI_latest)
 
 # create new aliases for io and ui

comfy_api/latest/_input/basic_types.py

@@ -1,5 +1,5 @@
 import torch
-from typing import TypedDict, List, Optional
+from typing import TypedDict, Optional
 
 ImageInput = torch.Tensor
 """
@@ -39,4 +39,4 @@ class LatentInput(TypedDict):
     Optional noise mask tensor in the same format as samples.
     """
 
-    batch_index: Optional[List[int]]
+    batch_index: Optional[list[int]]

comfy_api/latest/_ui.py

@@ -5,7 +5,6 @@ import os
 import random
 import uuid
 from io import BytesIO
-from typing import Type
 
 import av
 import numpy as np
@@ -83,7 +82,7 @@ class ImageSaveHelper:
         return PILImage.fromarray(np.clip(255.0 * image_tensor.cpu().numpy(), 0, 255).astype(np.uint8))
 
     @staticmethod
-    def _create_png_metadata(cls: Type[ComfyNode] | None) -> PngInfo | None:
+    def _create_png_metadata(cls: type[ComfyNode] | None) -> PngInfo | None:
         """Creates a PngInfo object with prompt and extra_pnginfo."""
         if args.disable_metadata or cls is None or not cls.hidden:
             return None
@@ -96,7 +95,7 @@ class ImageSaveHelper:
         return metadata
 
     @staticmethod
-    def _create_animated_png_metadata(cls: Type[ComfyNode] | None) -> PngInfo | None:
+    def _create_animated_png_metadata(cls: type[ComfyNode] | None) -> PngInfo | None:
         """Creates a PngInfo object with prompt and extra_pnginfo for animated PNGs (APNG)."""
         if args.disable_metadata or cls is None or not cls.hidden:
             return None
@@ -121,7 +120,7 @@ class ImageSaveHelper:
         return metadata
 
     @staticmethod
-    def _create_webp_metadata(pil_image: PILImage.Image, cls: Type[ComfyNode] | None) -> PILImage.Exif:
+    def _create_webp_metadata(pil_image: PILImage.Image, cls: type[ComfyNode] | None) -> PILImage.Exif:
         """Creates EXIF metadata bytes for WebP images."""
         exif_data = pil_image.getexif()
         if args.disable_metadata or cls is None or cls.hidden is None:
@@ -137,7 +136,7 @@ class ImageSaveHelper:
 
     @staticmethod
     def save_images(
-        images, filename_prefix: str, folder_type: FolderType, cls: Type[ComfyNode] | None, compress_level = 4,
+        images, filename_prefix: str, folder_type: FolderType, cls: type[ComfyNode] | None, compress_level = 4,
     ) -> list[SavedResult]:
         """Saves a batch of images as individual PNG files."""
         full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
@@ -155,7 +154,7 @@ class ImageSaveHelper:
         return results
 
     @staticmethod
-    def get_save_images_ui(images, filename_prefix: str, cls: Type[ComfyNode] | None, compress_level=4) -> SavedImages:
+    def get_save_images_ui(images, filename_prefix: str, cls: type[ComfyNode] | None, compress_level=4) -> SavedImages:
         """Saves a batch of images and returns a UI object for the node output."""
         return SavedImages(
                 ImageSaveHelper.save_images(
@@ -169,7 +168,7 @@ class ImageSaveHelper:
 
     @staticmethod
     def save_animated_png(
-        images, filename_prefix: str, folder_type: FolderType, cls: Type[ComfyNode] | None, fps: float, compress_level: int
+        images, filename_prefix: str, folder_type: FolderType, cls: type[ComfyNode] | None, fps: float, compress_level: int
     ) -> SavedResult:
         """Saves a batch of images as a single animated PNG."""
         full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
@@ -191,7 +190,7 @@ class ImageSaveHelper:
 
     @staticmethod
     def get_save_animated_png_ui(
-        images, filename_prefix: str, cls: Type[ComfyNode] | None, fps: float, compress_level: int
+        images, filename_prefix: str, cls: type[ComfyNode] | None, fps: float, compress_level: int
     ) -> SavedImages:
         """Saves an animated PNG and returns a UI object for the node output."""
         result = ImageSaveHelper.save_animated_png(
@@ -209,7 +208,7 @@ class ImageSaveHelper:
         images,
         filename_prefix: str,
         folder_type: FolderType,
-        cls: Type[ComfyNode] | None,
+        cls: type[ComfyNode] | None,
         fps: float,
         lossless: bool,
         quality: int,
@@ -238,7 +237,7 @@ class ImageSaveHelper:
     def get_save_animated_webp_ui(
         images,
         filename_prefix: str,
-        cls: Type[ComfyNode] | None,
+        cls: type[ComfyNode] | None,
         fps: float,
         lossless: bool,
         quality: int,
@@ -267,7 +266,7 @@ class AudioSaveHelper:
         audio: dict,
         filename_prefix: str,
         folder_type: FolderType,
-        cls: Type[ComfyNode] | None,
+        cls: type[ComfyNode] | None,
         format: str = "flac",
         quality: str = "128k",
     ) -> list[SavedResult]:
@@ -372,7 +371,7 @@ class AudioSaveHelper:
 
     @staticmethod
     def get_save_audio_ui(
-        audio, filename_prefix: str, cls: Type[ComfyNode] | None, format: str = "flac", quality: str = "128k",
+        audio, filename_prefix: str, cls: type[ComfyNode] | None, format: str = "flac", quality: str = "128k",
     ) -> SavedAudios:
         """Save and instantly wrap for UI."""
         return SavedAudios(
@@ -388,7 +387,7 @@ class AudioSaveHelper:
 
 
 class PreviewImage(_UIOutput):
-    def __init__(self, image: Image.Type, animated: bool = False, cls: Type[ComfyNode] = None, **kwargs):
+    def __init__(self, image: Image.Type, animated: bool = False, cls: type[ComfyNode] = None, **kwargs):
         self.values = ImageSaveHelper.save_images(
             image,
             filename_prefix="ComfyUI_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5)),
@@ -412,7 +411,7 @@ class PreviewMask(PreviewImage):
 
 
 class PreviewAudio(_UIOutput):
-    def __init__(self, audio: dict, cls: Type[ComfyNode] = None, **kwargs):
+    def __init__(self, audio: dict, cls: type[ComfyNode] = None, **kwargs):
         self.values = AudioSaveHelper.save_audio(
             audio,
             filename_prefix="ComfyUI_temp_" + "".join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5)),

comfy_api/version_list.py

@@ -2,9 +2,8 @@ from comfy_api.latest import ComfyAPI_latest
 from comfy_api.v0_0_2 import ComfyAPIAdapter_v0_0_2
 from comfy_api.v0_0_1 import ComfyAPIAdapter_v0_0_1
 from comfy_api.internal import ComfyAPIBase
-from typing import List, Type
 
-supported_versions: List[Type[ComfyAPIBase]] = [
+supported_versions: list[type[ComfyAPIBase]] = [
     ComfyAPI_latest,
     ComfyAPIAdapter_v0_0_2,
     ComfyAPIAdapter_v0_0_1,

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,