Initial support for RF models

(Hopefully) better integration with other node packs

Some internal cleanups

README.md

@@ -16,9 +16,10 @@ See the [changelog](changelog.md) for recent user visible changes.
 * There will be frequent workflow-breaking changes for a while yet.
 * Progress and previews are pretty wonky (you can look at the log for some progress information).
 * Using VAE or upscale models may result in the main model getting repeatedly unloaded/reloaded. Try using `latent` as the `guidance_mode`. If you actually have enough VRAM, maybe disabling smart memory (via ComfyUI commandline parameter) would help.
-* Currently only tested on SD15 and SDXL, may not work with models like Flux. (Not much testing in general as of yet.)
 * Brownian noise-based (AKA SDE) samplers may be a bit weird here, there is a workaround in place but it might not be enough. Also don't use with prompt-control's PCSplitSampling stuff.
 
+**Rectified Flow models note**: Should now work with RF models. SD3.5 apparently cannot handle high res images (even img2img) at all, so I don't recommend trying that. Flux seems to work pretty well. `image` guidance mode seems noticeably better than `latent` for Flux (based on my very limited testing) although it is slow. I haven't tested SD3.0 or other RF models, jank DiffuseHigh should handle them correctly but whether the results are actually decent I really couldn't say.
+
 ## Description
 
 The DiffuseHigh approach is similar to an iterative upscale/run some more steps at low denoise approach with a twist: it mixes in guidance from a reference image for a number of steps at the beginning of each sampling iteration. The guidance is derived from the low frequency parts of the reference and also gets sharpened first to increase detail.
@@ -235,6 +236,7 @@ I tried to set the node defaults to align with the official implementation. Thes
 * It's very important that the initial reference is as close to flawless as possible. Unlike the normal highres fix approach which can sometimes fix issues when you set the denoise relatively high, DiffuseHigh guidance keeps the model from diverging from the reference too much. This can be a double edged sword in some cases.
 * The sampler has a workaround for a [long standing bug in ComfyUI](https://github.com/comfyanonymous/ComfyUI/issues/2833) where generations aren't deterministic when `add_noise` is disabled in the sampler. However, this may change seeds. You can disable the workaround via the advanced YAML options - see `seed_rng` and `seed_rng_offset`.
 * For `taesd` VAE mode, you will need the TAESD encoder models available at https://github.com/madebyollin/taesd - put them in `models/vae_approx`.
+* You can use DiffuseHigh as an enhanced highres-fix by passing a pre-upscaled reference image, setting the iteration count to one and using a scale factor of 1.0.
 
 ***
 

changelog.md

@@ -2,6 +2,10 @@
 
 Note, only relatively significant changes to user-visible functionality will be included here. Most recent changes at the top.
 
+## 20241023
+
+* Initial support for rectified flow models (Flux, SD3, SD3.5). Might slightly change seeds for other models.
+
 ## 20241019
 
 * Added workaround for https://github.com/comfyanonymous/ComfyUI/issues/2833 - may change seeds. You can disable it with `seed_rng: false` in YAML parameters.

py/external.py

@@ -1,16 +1,46 @@
 import contextlib
-import importlib
+import importlib.util
+import sys
 
 EXTERNAL = {}
+INITIALIZED = False
 
-with contextlib.suppress(ImportError):
-    EXTERNAL["tiled_diffusion"] = importlib.import_module(
-        "custom_nodes.ComfyUI-TiledDiffusion",
-    )
 
-with contextlib.suppress(ImportError, NotImplementedError):
-    bleh = importlib.import_module("custom_nodes.ComfyUI-bleh")
-    bleh_version = getattr(bleh, "BLEH_VERSION", -1)
-    if bleh_version < 1:
-        raise NotImplementedError
-    EXTERNAL["bleh"] = bleh.py
+def get_custom_node(name):
+    module_key = f"custom_nodes.{name}"
+    try:
+        spec = importlib.util.find_spec(module_key)
+        if spec is None:
+            raise ModuleNotFoundError(module_key)
+        module = next(
+            v
+            for v in sys.modules.copy().values()
+            if hasattr(v, "__spec__")
+            and v.__spec__ is not None
+            and v.__spec__.origin == spec.origin
+        )
+    except StopIteration:
+        raise ModuleNotFoundError(module_key) from None
+    return module
+
+
+def init_integrations() -> None:
+    global INITIALIZED  # noqa: PLW0603
+    if INITIALIZED:
+        return
+    INITIALIZED = True
+
+    with contextlib.suppress(ModuleNotFoundError):
+        EXTERNAL["tiled_diffusion"] = get_custom_node("ComfyUI-TiledDiffusion")
+
+    with contextlib.suppress(ModuleNotFoundError, NotImplementedError):
+        bleh = get_custom_node("ComfyUI-bleh")
+        bleh_version = getattr(bleh, "BLEH_VERSION", -1)
+        if bleh_version < 1:
+            raise NotImplementedError
+        EXTERNAL["bleh"] = bleh.py
+
+    from . import tensor_image_ops, vae  # noqa: PLC0415
+
+    tensor_image_ops.init_integrations()
+    vae.init_integrations()

py/nodes.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import yaml
 from comfy.samplers import KSAMPLER
 
+from .external import init_integrations
 from .sampler import diffusehigh_sampler
 from .vae import VAEMode
 
@@ -141,6 +142,7 @@ class DiffuseHighSamplerNode:
 
     @classmethod
     def go(cls, yaml_parameters: None | str = None, **kwargs: dict) -> tuple[KSAMPLER]:
+        init_integrations()
         if yaml_parameters:
             extra_params = yaml.safe_load(yaml_parameters)
             if extra_params is None:

py/sampler.py

@@ -300,6 +300,10 @@ class DiffuseHighSampler:
     def __call__(self) -> torch.Tensor:
         self.config = self.base_config.get_iteration_config("reference")
         if self.reference_image is None:
+            if self.sigmas[-1] != 0:
+                raise ValueError(
+                    "Initial reference sigmas must end at 0 (full denoise)",
+                )
             normal_step_count = len(self.sigmas) - 1
             with tqdm(
                 disable=self.disable_pbar,
@@ -327,6 +331,10 @@ class DiffuseHighSampler:
             del x_new
             self.config = self.base_config.get_iteration_config(iteration)
             self.set_schedule()
+            if self.highres_sigmas[-1] != 0:
+                raise ValueError(
+                    "Highres sigmas (including schedule overrides) must end at 0 (full denoise)",
+                )
             self.gc()
             if self.config.sigma_offset >= len(self.highres_sigmas) - 1:
                 raise ValueError(
@@ -367,15 +375,14 @@ class DiffuseHighSampler:
                 raise ValueError("Bad guidance_mode")
             x_noise = torch.randn_like(x_new)
             self.seed_offset += 1
-            x_new = x_new + x_noise * (
-                self.highres_sigmas[self.sigma_offset] * self.renoise_factor
+            x_new = self.model_sampling.noise_scaling(
+                self.highres_sigmas[self.sigma_offset],
+                x_noise.mul_(self.renoise_factor),
+                x_new,
+                max_denoise=self.highres_sigmas[self.sigma_offset]
+                >= self.model_sampling.sigma_max - 1e-05,
             )
             del x_noise
-            # x_new = self.model.inner_model.inner_model.model_sampling.noise_scaling(
-            #     self.highres_sigmas[0] * self.renoise_factor,
-            #     x_noise,
-            #     x_new,
-            # )
             self.gc()
             x_new = self.run_steps(x_new, sigmas=self.highres_sigmas)
             if iteration == self.iterations - 1:

py/tensor_image_ops.py

@@ -12,14 +12,17 @@ from .external import EXTERNAL
 
 F = torch.nn.functional
 
-EXT_BLEH = EXTERNAL.get("bleh")
+BLENDING_MODES = {
+    "lerp": torch.lerp,
+}
 
-if EXT_BLEH is not None:
-    BLENDING_MODES = EXT_BLEH.latent_utils.BLENDING_MODES
-else:
-    BLENDING_MODES = {
-        "lerp": torch.lerp,
-    }
+
+def init_integrations():
+    global BLENDING_MODES  # noqa: PLW0603
+    ext_bleh = EXTERNAL.get("bleh")
+    if ext_bleh is not None:
+        BLENDING_MODES.clear()
+        BLENDING_MODES |= ext_bleh.latent_utils.BLENDING_MODES
 
 
 class SharpenMode(Enum):

py/vae.py

@@ -10,7 +10,12 @@ from tqdm import tqdm
 from .external import EXTERNAL
 from .utils import fallback
 
-tiled_diffusion = EXTERNAL.get("tiled_diffusion")
+tiled_diffusion = None
+
+
+def init_integrations():
+    global tiled_diffusion  # noqa: PLW0603
+    tiled_diffusion = EXTERNAL.get("tiled_diffusion")
 
 
 class VAEMode(Enum):
@@ -84,25 +89,20 @@ class VAEHelper:
     def encode(self, imgbatch, *, fix_dims=False, disable_pbar=None):
         if fix_dims:
             imgbatch = imgbatch.moveaxis(1, -1)
-        # print("ENCODING", imgbatch.min(), imgbatch.max())
         with tqdm(disable=disable_pbar, total=1, desc="VAE encode") as pbar:
             result = self.encode_fun(imgbatch[..., :3])
             pbar.update()
         if self.mode != VAEMode.TAESD:
-            # print("ENCODED(raw):", result.min(), result.max())
             result = self.latent_format.process_in(result)
-        # print("ENCODED", result.shape, result.min(), result.max())
         return result
 
     def decode(self, latent, *, skip_process_out=False, disable_pbar=None):
         if self.mode != VAEMode.TAESD and not skip_process_out:
             latent = self.latent_format.process_out(latent)
-        # print("DECODING", latent.min(), latent.max())
         with tqdm(disable=disable_pbar, total=1, desc="VAE decode") as pbar:
             result = self.decode_fun(latent)
             pbar.update()
         return result
-        # print("DECODED", result.shape, result.min(), result.max())
 
     def encode_taesd(self, imgbatch):
         dummy = torch.zeros((), device=self.device, dtype=self.dtype)
@@ -112,7 +112,6 @@ class VAEHelper:
         return OCSTAESD.decode(self.latent_format, latent)
 
     def encode_vae(self, imgbatch):
-        # print("VAE ENC", imgbatch.shape)
         return self.vae.encode(imgbatch, **self.encode_kwargs)
 
     def decode_vae(self, latent):