Add multi-GPU device routing

README.md

@@ -140,6 +140,72 @@ ComfyUI:
 export HF_HOME=/path/to/huggingface/cache
 ```
 
+## Experimental split-device inference
+
+The pipeline loader supports separate diffusion and text-encoder devices
+without changing existing workflows. Set these environment variables before
+starting ComfyUI:
+
+```bash
+export IDEOGRAM4_DIFFUSION_DEVICE=cuda:0
+export IDEOGRAM4_TEXT_DEVICE=cuda:1
+```
+
+`IDEOGRAM4_DIFFUSION_DEVICE` owns both diffusion transformers, the VAE, sampling
+state, and decode state. `IDEOGRAM4_TEXT_DEVICE` owns Qwen3-VL. Qwen runs once
+per prompt, then compact text-only features move once to the diffusion device
+before denoising. If either variable is unset, it defaults to `cuda`; if only
+the text variable is unset, it inherits the diffusion device, preserving the
+single-device behavior.
+
+AMD ROCm multi-GPU use is experimental. ROCm PyTorch exposes Radeon devices
+through `torch.cuda`, so use `cuda:N` names and check the runtime mapping first:
+
+```bash
+python - <<'PY'
+import torch
+
+print("torch:", torch.__version__)
+print("HIP:", torch.version.hip)
+print("device count:", torch.cuda.device_count())
+for i in range(torch.cuda.device_count()):
+    properties = torch.cuda.get_device_properties(i)
+    print(
+        f"{i}: {torch.cuda.get_device_name(i)} "
+        f"VRAM={properties.total_memory / 1024**3:.2f} GiB"
+    )
+if torch.cuda.device_count() >= 2:
+    print("0 -> 1 peer:", torch.cuda.can_device_access_peer(0, 1))
+    print("1 -> 0 peer:", torch.cuda.can_device_access_peer(1, 0))
+PY
+```
+
+Install the modified core checkout into the same environment as ComfyUI and
+verify its import path:
+
+```bash
+python -m pip uninstall -y ideogram-4
+python -m pip install -e /absolute/path/to/ideogram4
+python - <<'PY'
+import ideogram4
+print(ideogram4.__file__)
+PY
+```
+
+Example ROCm launch:
+
+```bash
+export ROCR_VISIBLE_DEVICES=0,1
+export IDEOGRAM4_REPO=/absolute/path/to/ideogram4
+export IDEOGRAM4_DIFFUSION_DEVICE=cuda:0
+export IDEOGRAM4_TEXT_DEVICE=cuda:1
+python main.py --listen 0.0.0.0 --port 8189 --disable-pinned-memory
+```
+
+Use explicit indexed devices for split placement. Direct peer access is not
+required; the core pipeline warns and stages the once-per-prompt feature copy
+through CPU if a direct transfer fails.
+
 ## Quick Workflow
 
 1. Add `Ideogram 4.0 Magic Prompt`.

nodes.py

@@ -29,7 +29,8 @@
 
 DEFAULT_MODEL_WEIGHTS = "4.0 NF4"
 DEFAULT_TORCH_DTYPE = torch.bfloat16
-DEFAULT_DEVICE = "cuda"
+DEFAULT_DEVICE = os.environ.get("IDEOGRAM4_DIFFUSION_DEVICE", "cuda")
+DEFAULT_TEXT_DEVICE = os.environ.get("IDEOGRAM4_TEXT_DEVICE", DEFAULT_DEVICE)
 
 CUSTOM_SAMPLER_PRESET = "custom"
 MAGIC_PROMPT_PROVIDER_IDEOGRAM = "ideogram"
@@ -524,6 +525,8 @@ def load(
       model_weights,
       resolved_weights_repo,
       os.environ.get(CORE_REPO_ENV_VAR, ""),
+      DEFAULT_DEVICE,
+      DEFAULT_TEXT_DEVICE,
     )
     if cache_key in self._CACHE:
       _send_progress_text(unique_id, _status_text("Status: Ready from cache", f"Weights: {model_weights}"))
@@ -542,10 +545,16 @@ def load(
       weights_repo=resolved_weights_repo,
     )
     pbar.update_absolute(2)
+    print(
+      "Ideogram 4.0 Pipeline Loader devices: "
+      f"diffusion_device={DEFAULT_DEVICE}, text_device={DEFAULT_TEXT_DEVICE}",
+      flush=True,
+    )
     try:
       pipeline = Ideogram4Pipeline.from_pretrained(
         config=config,
         device=DEFAULT_DEVICE,
+        text_device=DEFAULT_TEXT_DEVICE,
         dtype=DEFAULT_TORCH_DTYPE,
       )
     except Exception as exc:

tests/test_device_routing.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import importlib.util
+import os
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest import mock
+
+
+class _ProgressBar:
+  def update_absolute(self, _value):
+    pass
+
+
+def _load_nodes_module():
+  comfy = types.ModuleType("comfy")
+  comfy_utils = types.ModuleType("comfy.utils")
+  comfy_utils.ProgressBar = _ProgressBar
+  comfy.utils = comfy_utils
+
+  module_path = Path(__file__).parents[1] / "nodes.py"
+  spec = importlib.util.spec_from_file_location(
+    "ideogram4_comfy_test_nodes", module_path
+  )
+  if spec is None or spec.loader is None:
+    raise RuntimeError(f"Unable to load {module_path}")
+  module = importlib.util.module_from_spec(spec)
+  with (
+    mock.patch.dict(
+      os.environ,
+      {
+        "IDEOGRAM4_DIFFUSION_DEVICE": "cuda:3",
+        "IDEOGRAM4_TEXT_DEVICE": "cuda:7",
+      },
+    ),
+    mock.patch.dict(
+      sys.modules,
+      {"comfy": comfy, "comfy.utils": comfy_utils},
+    ),
+  ):
+    spec.loader.exec_module(module)
+  return module
+
+
+class DeviceRoutingTests(unittest.TestCase):
+  def test_loader_passes_devices_and_separates_cache_key(self) -> None:
+    nodes = _load_nodes_module()
+    captured: dict[str, object] = {}
+
+    class FakeConfig:
+      def __init__(self, **kwargs):
+        captured["config"] = kwargs
+
+    class FakePipeline:
+      @classmethod
+      def from_pretrained(cls, **kwargs):
+        captured["pipeline_kwargs"] = kwargs
+        return object()
+
+    loader = nodes.Ideogram4PipelineLoader()
+    loader._CACHE = {}
+    with (
+      mock.patch.object(nodes.torch.cuda, "is_available", return_value=True),
+      mock.patch.object(nodes, "_apply_hf_token"),
+      mock.patch.object(
+        nodes,
+        "_load_pipeline_classes",
+        return_value=(FakePipeline, FakeConfig),
+      ),
+      mock.patch.object(nodes, "_progress", return_value=_ProgressBar()),
+      mock.patch.object(nodes, "_send_progress_text"),
+    ):
+      loader.load("4.0 NF4")
+
+    pipeline_kwargs = captured["pipeline_kwargs"]
+    self.assertEqual(pipeline_kwargs["device"], "cuda:3")
+    self.assertEqual(pipeline_kwargs["text_device"], "cuda:7")
+    self.assertIn(
+      (
+        "4.0 NF4",
+        "ideogram-ai/ideogram-4-nf4",
+        "",
+        "cuda:3",
+        "cuda:7",
+      ),
+      loader._CACHE,
+    )
+
+
+if __name__ == "__main__":
+  unittest.main()