avtc
01a324bb11
Extends tiling beyond UNet (Conv2d) models to support DiT architectures using toroidal attention patches and latent content wrapping for seamless infinite tiling. Adds rectangular tiling mode, WanVAE 5D tensor handling, and an AdvancedTilingRay node for distributed Raylight workers.
82 lines
2.8 KiB
Python
82 lines
2.8 KiB
Python
"""
|
|
Raylight integration for DiT tiling.
|
|
|
|
Provides AdvancedTilingRay node that works with raylight's RAY_ACTORS type,
|
|
applying toroidal attention tiling on distributed workers.
|
|
|
|
Uses manual Ray remote calls instead of @ray_patch to avoid serialization
|
|
issues with custom node packages that have non-standard import paths
|
|
(e.g. hyphens in directory names that prevent normal Python imports).
|
|
"""
|
|
|
|
try:
|
|
from raylight.comfy_extra_dist.ray_patch_decorator import ray_patch
|
|
HAS_RAYLIGHT = True
|
|
except ImportError:
|
|
HAS_RAYLIGHT = False
|
|
|
|
if HAS_RAYLIGHT:
|
|
import ray
|
|
|
|
class AdvancedTilingRay:
|
|
"""
|
|
Applies tiling to a raylight RAY_ACTORS model.
|
|
Dispatches per-worker patching via Ray remote calls.
|
|
"""
|
|
|
|
# pylint: disable=invalid-name
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"settings": ("ADVANCED_TILING_SETTINGS",),
|
|
"ray_actors": ("RAY_ACTORS",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("RAY_ACTORS",)
|
|
RETURN_NAMES = ("ray_actors",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "conditioning"
|
|
|
|
def run(self, settings, ray_actors):
|
|
# Defined inside run() so cloudpickle serializes it as a nested
|
|
# function (by value) instead of by module reference. Module-level
|
|
# functions get serialized by reference, which requires importing
|
|
# the module by name on the worker -- but the module name is the
|
|
# filesystem path (with a hyphen), causing ModuleNotFoundError.
|
|
def _patch(model, mode, rotation):
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
|
|
_node_dir = os.path.dirname(os.path.abspath(__file__))
|
|
_pkg = "ComfyUI_AdvancedTiling"
|
|
|
|
if _pkg not in sys.modules:
|
|
_init = os.path.join(_node_dir, "__init__.py")
|
|
spec = importlib.util.spec_from_file_location(
|
|
_pkg, _init,
|
|
submodule_search_locations=[_node_dir],
|
|
)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
sys.modules[_pkg] = mod
|
|
spec.loader.exec_module(mod)
|
|
|
|
from ComfyUI_AdvancedTiling.dit_tiling import patch_dit_model
|
|
from ComfyUI_AdvancedTiling.modes import Settings
|
|
|
|
patch_dit_model(model, Settings(mode, rotation))
|
|
return model
|
|
|
|
gpu_workers = ray_actors["workers"]
|
|
futures = [
|
|
actor.model_function_runner.remote(
|
|
_patch, settings.mode, settings.rotation
|
|
)
|
|
for actor in gpu_workers
|
|
]
|
|
ray.get(futures)
|
|
return (ray_actors,)
|