Add tiled VAE lane to MultiGPU Work Units

tests-unit/comfy_test/multigpu_test.py

@@ -0,0 +1,147 @@
+import importlib
+import sys
+import types
+
+import torch
+
+import comfy.utils
+
+
+def install_fake_comfy_aimdo(monkeypatch):
+    package = types.ModuleType("comfy_aimdo")
+    package.__path__ = []
+    monkeypatch.setitem(sys.modules, "comfy_aimdo", package)
+    for name in ("vram_buffer", "host_buffer", "torch", "model_vbar", "model_mmap", "control"):
+        module = types.ModuleType(f"comfy_aimdo.{name}")
+        monkeypatch.setitem(sys.modules, f"comfy_aimdo.{name}", module)
+        setattr(package, name, module)
+
+
+def test_tiled_scale_multidim_multigpu_clips_edge_tiles(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "set_device", lambda device: None)
+    monkeypatch.setattr(torch.cuda, "synchronize", lambda device: None)
+
+    scale = 1.1
+
+    def upscale(a):
+        return torch.ones((a.shape[0], 1, round(a.shape[-1] * scale)), dtype=a.dtype, device=a.device)
+
+    samples = torch.ones((1, 1, 11))
+    devices = [torch.device("cpu:0"), torch.device("cpu:1")]
+
+    actual = comfy.utils.tiled_scale_multidim_multigpu(
+        samples,
+        {device: upscale for device in devices},
+        tile=(5,),
+        overlap=2,
+        upscale_amount=scale,
+        out_channels=1,
+        output_device="cpu",
+    )
+    expected = comfy.utils.tiled_scale_multidim(
+        samples,
+        upscale,
+        tile=(5,),
+        overlap=2,
+        upscale_amount=scale,
+        out_channels=1,
+        output_device="cpu",
+    )
+
+    assert actual.shape == expected.shape == (1, 1, 12)
+    torch.testing.assert_close(actual, expected)
+
+
+def test_upscale_model_deepclone_does_not_copy_existing_clone_graph(monkeypatch):
+    class FakeModel:
+        def __init__(self):
+            self.param = torch.nn.Parameter(torch.ones(1))
+
+        def eval(self):
+            return self
+
+        def parameters(self):
+            return [self.param]
+
+    class FakeDescriptor:
+        def __init__(self):
+            self.model = FakeModel()
+            self.device = None
+
+        def to(self, device):
+            self.device = device
+            return self
+
+    first_device = torch.device("cpu:0")
+    second_device = torch.device("cpu:1")
+    stale_device = torch.device("cpu:2")
+    existing_clone = FakeDescriptor()
+    stale_clone = FakeDescriptor()
+    source = FakeDescriptor()
+    source.multigpu_clones = {first_device: existing_clone, stale_device: stale_clone}
+    fake_model_management = types.ModuleType("comfy.model_management")
+    fake_model_management.get_all_torch_devices = lambda exclude_current=True: [first_device, second_device]
+    monkeypatch.setitem(sys.modules, "comfy.model_management", fake_model_management)
+    import comfy
+    monkeypatch.setattr(comfy, "model_management", fake_model_management, raising=False)
+    import comfy.multigpu
+    importlib.reload(comfy.multigpu)
+
+    cloned = comfy.multigpu.create_upscale_model_multigpu_deepclones(source, max_gpus=3)
+
+    assert cloned is not source
+    assert cloned.multigpu_clones[first_device] is existing_clone
+    assert stale_device not in cloned.multigpu_clones
+    assert second_device in cloned.multigpu_clones
+    assert not hasattr(cloned.multigpu_clones[second_device], "multigpu_clones")
+    assert cloned.multigpu_clones[second_device].device == "cpu"
+    assert not cloned.multigpu_clones[second_device].model.param.requires_grad
+
+    single_gpu_clone = comfy.multigpu.create_upscale_model_multigpu_deepclones(source, max_gpus=1)
+    assert single_gpu_clone is not source
+    assert not hasattr(single_gpu_clone, "multigpu_clones")
+
+
+def test_checkpoint_loader_registers_vae_cached_patcher(monkeypatch):
+    install_fake_comfy_aimdo(monkeypatch)
+    import comfy.sd
+    importlib.reload(comfy.sd)
+
+    class FakeVAE:
+        def __init__(self):
+            self.patcher = types.SimpleNamespace(cached_patcher_init=None)
+
+    model_patcher = types.SimpleNamespace(cached_patcher_init=None)
+    vae = FakeVAE()
+    metadata = {"format": "checkpoint"}
+    monkeypatch.setattr(comfy.utils, "load_torch_file", lambda path, return_metadata=False: ({}, metadata))
+    monkeypatch.setattr(
+        comfy.sd,
+        "load_state_dict_guess_config",
+        lambda *args, **kwargs: (model_patcher, None, vae, None),
+    )
+
+    comfy.sd.load_checkpoint_guess_config("checkpoint.safetensors", output_vae=True)
+
+    assert model_patcher.cached_patcher_init[0] is comfy.sd.load_checkpoint_guess_config
+    assert vae.patcher.cached_patcher_init[0] is comfy.sd.load_checkpoint_vae_patcher
+    assert vae.patcher.cached_patcher_init[1][0] == "checkpoint.safetensors"
+
+
+def test_checkpoint_loader_skips_cached_patcher_for_placeholder_vae(monkeypatch):
+    install_fake_comfy_aimdo(monkeypatch)
+    import comfy.sd
+    importlib.reload(comfy.sd)
+
+    model_patcher = types.SimpleNamespace(cached_patcher_init=None)
+    placeholder_vae = types.SimpleNamespace()
+    metadata = {"format": "checkpoint"}
+    monkeypatch.setattr(comfy.utils, "load_torch_file", lambda path, return_metadata=False: ({}, metadata))
+    monkeypatch.setattr(
+        comfy.sd,
+        "load_state_dict_guess_config",
+        lambda *args, **kwargs: (model_patcher, None, placeholder_vae, None),
+    )
+
+    assert comfy.sd.load_checkpoint_guess_config("diffusion_only.safetensors", output_vae=True)[2] is placeholder_vae
+    assert model_patcher.cached_patcher_init[0] is comfy.sd.load_checkpoint_guess_config