From af8f68f38c600e73e3572ca807f7fe948bd1fc3e Mon Sep 17 00:00:00 2001 From: jason Date: Fri, 31 Jul 2026 14:19:50 -0700 Subject: [PATCH 1/2] Route Ideogram pipeline devices from environment --- nodes.py | 11 ++++- tests/test_device_routing.py | 93 ++++++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 1 deletion(-) create mode 100644 tests/test_device_routing.py diff --git a/nodes.py b/nodes.py index 8510b8c..f5ace34 100644 --- a/nodes.py +++ b/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: diff --git a/tests/test_device_routing.py b/tests/test_device_routing.py new file mode 100644 index 0000000..60c78a4 --- /dev/null +++ b/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() From 1b70f9c50d077a35c26a8ae5c16ede0fc514aa60 Mon Sep 17 00:00:00 2001 From: jason Date: Fri, 31 Jul 2026 14:19:51 -0700 Subject: [PATCH 2/2] Document ComfyUI split-device setup --- README.md | 66 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/README.md b/README.md index 641a985..dc8b108 100644 --- a/README.md +++ b/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`.