diff --git a/README.md b/README.md index a2154ba..e5bf4ed 100644 --- a/README.md +++ b/README.md @@ -154,6 +154,95 @@ under `src/ideogram4/` take effect without reinstalling: pip install -e . ``` +### Experimental split-device inference + +`Ideogram4Pipeline.from_pretrained` accepts an optional `text_device`. This +keeps Qwen3-VL on that device while the conditional and unconditional diffusion +transformers, VAE, sampling state, and decode state remain on `device`. If +`text_device` is omitted, all components use `device` exactly as before. + +```python +import torch + +pipeline = Ideogram4Pipeline.from_pretrained( + device="cuda:0", + text_device="cuda:1", + dtype=torch.bfloat16, +) +``` + +On split devices, Qwen runs once over the left-padded text block only. The +resulting float32 text features are copied once to the diffusion device before +sampling, and image-position zeros are allocated there. There are no planned +cross-device transfers in the denoising loop. A direct GPU transfer is tried +first; runtimes without peer transfer support fall back to CPU staging with a +warning. + +AMD ROCm multi-GPU support is experimental. PyTorch ROCm still uses the +`cuda:N` device spelling. Use explicit indexes for split placement and verify +the runtime mapping rather than assuming device order: + +```bash +python scripts/smoke_split_devices.py \ + --caption-file /path/to/structured-caption.json \ + --output /tmp/ideogram4-smoke.png \ + --diffusion-device cuda:0 \ + --text-device cuda:1 \ + --height 256 --width 256 --num-steps 2 +``` + +The smoke script prints visible devices, peer-access status, and peak allocated +and reserved memory. It does not call a hosted Magic Prompt API. + +After the small smoke passes, run the acceptance configuration with the +official 20-step sampler schedule: + +```bash +python scripts/smoke_split_devices.py \ + --caption-file /path/to/structured-caption.json \ + --output /tmp/ideogram4-1024.png \ + --diffusion-device cuda:0 \ + --text-device cuda:1 \ + --height 1024 --width 1024 \ + --sampler-preset V4_DEFAULT_20 +``` + +For the target Fedora/ROCm machine, the validation wrapper performs the full +preflight, installs this checkout editable without replacing the existing ROCm +PyTorch build, records `rocm-smi` telemetry, runs both core tests above, and can +then launch ComfyUI: + +From the ComfyUI directory, this one command safely clones or updates both +feature branches, creates a local structured test caption, runs the complete +validator, and launches ComfyUI: + +```bash +bash <(curl -fsSL https://raw.githubusercontent.com/Dillflix/ideogram4/feature/split-text-diffusion-devices/scripts/fedora_one_step.sh) +``` + +If ComfyUI is elsewhere, identify it on the same command: + +```bash +COMFYUI_DIR=/absolute/path/to/ComfyUI bash <(curl -fsSL https://raw.githubusercontent.com/Dillflix/ideogram4/feature/split-text-diffusion-devices/scripts/fedora_one_step.sh) +``` + +The lower-level equivalent, useful when the repositories are already prepared, +is: + +```bash +chmod +x scripts/validate_rocm_split_devices.sh +scripts/validate_rocm_split_devices.sh \ + --caption /path/to/structured-caption.json \ + --comfyui-dir /absolute/path/to/ComfyUI \ + --launch-comfyui +``` + +The script fails before loading weights if the resolved device names do not +contain `7900 XT` for the diffusion role and `8060S` for the text role. Use its +explicit device/name options if the runtime mapping differs. Logs, peak-memory +reports, generated images, and GPU telemetry are written to a timestamped +validation directory. + ### Model access The model weights are **gated** on Hugging Face, so you must accept the gate and diff --git a/pyproject.toml b/pyproject.toml index 00bcd96..7c0b019 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,10 @@ packages = ["src/ideogram4"] include = [ "src/ideogram4", "src/ideogram4/magic_prompt_system_prompts/*.txt", + "scripts/fedora_one_step.sh", + "scripts/smoke_split_devices.py", + "scripts/validate_rocm_split_devices.sh", + "tests", "run_inference.py", "README.md", "pyproject.toml", diff --git a/scripts/fedora_one_step.sh b/scripts/fedora_one_step.sh new file mode 100755 index 0000000..123fa89 --- /dev/null +++ b/scripts/fedora_one_step.sh @@ -0,0 +1,167 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +CORE_REPO="https://github.com/Dillflix/ideogram4.git" +CORE_BRANCH="feature/split-text-diffusion-devices" +WRAPPER_REPO="https://github.com/Dillflix/ComfyUI-Ideogram4.git" +WRAPPER_BRANCH="feature/multigpu-device-routing" +COMFYUI_REPO="https://github.com/comfyanonymous/ComfyUI.git" +REMOTE_NAME="dillflix-validation" +STATE_ROOT="${XDG_DATA_HOME:-$HOME/.local/share}/ideogram4-split-validation" +CORE_DIR="$STATE_ROOT/ideogram4" + +fail() { + echo "ERROR: $*" >&2 + exit 1 +} + +is_comfyui_root() { + [[ -f "$1/main.py" && -f "$1/nodes.py" && -d "$1/comfy" ]] +} + +find_comfyui() { + local candidate venv_parent + local -a candidates=() + if [[ -n "${COMFYUI_DIR:-}" ]]; then + candidate="$COMFYUI_DIR" + is_comfyui_root "$candidate" || fail "COMFYUI_DIR is not a ComfyUI checkout: $candidate" + readlink -f -- "$candidate" + return + fi + + if [[ -n "${VIRTUAL_ENV:-}" ]]; then + venv_parent="$(dirname -- "$VIRTUAL_ENV")" + candidates+=("$venv_parent" "$(dirname -- "$venv_parent")") + fi + candidates+=("$PWD" "$HOME/ComfyUI" "$HOME/comfyui") + + for candidate in "${candidates[@]}"; do + if is_comfyui_root "$candidate"; then + readlink -f -- "$candidate" + return + fi + done + + while IFS= read -r candidate; do + candidate="$(dirname -- "$candidate")" + if is_comfyui_root "$candidate"; then + echo "$candidate" + return + fi + done < <( + find "$HOME" -maxdepth 8 -type f -name main.py -ipath '*comfyui*' \ + -print 2>/dev/null | sort + ) + + candidate="$STATE_ROOT/ComfyUI" + echo "Existing ComfyUI checkout not found; cloning it into $candidate" >&2 + if [[ ! -d "$candidate/.git" ]]; then + mkdir -p -- "$(dirname -- "$candidate")" + git clone "$COMFYUI_REPO" "$candidate" >&2 + fi + readlink -f -- "$candidate" +} + +checkout_branch() { + local repo_url="$1" + local branch="$2" + local destination="$3" + + if [[ ! -d "$destination/.git" ]]; then + if [[ -e "$destination" && -n "$(find "$destination" -mindepth 1 -maxdepth 1 -print -quit 2>/dev/null)" ]]; then + fail "cannot clone into non-empty directory: $destination" + fi + mkdir -p -- "$(dirname -- "$destination")" + git clone --branch "$branch" --single-branch "$repo_url" "$destination" + return + fi + + if [[ -n "$(git -C "$destination" status --porcelain)" ]]; then + fail "existing checkout has local changes; refusing to switch it: $destination" + fi + + if git -C "$destination" remote get-url "$REMOTE_NAME" >/dev/null 2>&1; then + git -C "$destination" remote set-url "$REMOTE_NAME" "$repo_url" + else + git -C "$destination" remote add "$REMOTE_NAME" "$repo_url" + fi + git -C "$destination" fetch "$REMOTE_NAME" "$branch" + + if git -C "$destination" show-ref --verify --quiet "refs/heads/$branch"; then + git -C "$destination" switch "$branch" + else + git -C "$destination" switch --create "$branch" --track "$REMOTE_NAME/$branch" + fi + git -C "$destination" merge --ff-only "$REMOTE_NAME/$branch" +} + +command -v git >/dev/null 2>&1 || fail "git is required" +[[ "$(uname -s)" == "Linux" ]] || fail "this launcher must run on Linux" + +COMFYUI_DIR="$(find_comfyui)" +WRAPPER_DIR="$COMFYUI_DIR/custom_nodes/ComfyUI-Ideogram4" +mkdir -p -- "$STATE_ROOT" + +echo "ComfyUI: $COMFYUI_DIR" +echo "Validation workspace: $STATE_ROOT" +echo "Preparing modified core..." +checkout_branch "$CORE_REPO" "$CORE_BRANCH" "$CORE_DIR" +echo "Preparing modified ComfyUI wrapper..." +checkout_branch "$WRAPPER_REPO" "$WRAPPER_BRANCH" "$WRAPPER_DIR" + +CAPTION_FILE="$STATE_ROOT/validation-caption.json" +cat >"$CAPTION_FILE" <<'JSON' +{ + "high_level_description": "A clean technical poster celebrating a dual-GPU image generation workstation.", + "style_description": { + "aesthetics": "precise, modern, polished, high contrast", + "lighting": "soft studio lighting with subtle cyan and amber rim light", + "medium": "graphic_design", + "art_style": "minimal technical poster, crisp geometric forms, premium product visualization", + "color_palette": ["#101820", "#00AEEF", "#FFB000", "#F5F7FA"] + }, + "compositional_deconstruction": { + "background": "A deep charcoal studio backdrop with a faint grid and restrained cyan highlights.", + "elements": [ + { + "type": "obj", + "bbox": [180, 120, 820, 880], + "desc": "Two elegant abstract GPU modules connected by a single luminous data arc, arranged symmetrically as a premium technical product composition." + }, + { + "type": "text", + "bbox": [70, 180, 180, 820], + "text": "SPLIT COMPUTE", + "desc": "Large crisp uppercase geometric sans-serif title in white." + } + ] + } +} +JSON + +echo "Starting automated ROCm validation." +echo "This downloads gated NF4 weights if they are not already cached." +echo "Core and ComfyUI API tests run automatically; the validated server then remains running." + +VALIDATOR_ARGS=( + --caption "$CAPTION_FILE" + --comfyui-dir "$COMFYUI_DIR" + --launch-comfyui +) + +if [[ -n "${VIRTUAL_ENV:-}" && -x "$VIRTUAL_ENV/bin/python" ]]; then + VALIDATOR_ARGS+=(--python "$VIRTUAL_ENV/bin/python") +fi + +if [[ -n "${IDEOGRAM4_VALIDATION_OUTPUT:-}" ]]; then + VALIDATOR_ARGS+=(--output-dir "$IDEOGRAM4_VALIDATION_OUTPUT") +fi +if [[ -n "${IDEOGRAM4_DIFFUSION_DEVICE:-}" ]]; then + VALIDATOR_ARGS+=(--diffusion-device "$IDEOGRAM4_DIFFUSION_DEVICE") +fi +if [[ -n "${IDEOGRAM4_TEXT_DEVICE:-}" ]]; then + VALIDATOR_ARGS+=(--text-device "$IDEOGRAM4_TEXT_DEVICE") +fi + +exec "$CORE_DIR/scripts/validate_rocm_split_devices.sh" "${VALIDATOR_ARGS[@]}" diff --git a/scripts/smoke_split_devices.py b/scripts/smoke_split_devices.py new file mode 100644 index 0000000..ceb7def --- /dev/null +++ b/scripts/smoke_split_devices.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +import torch + +from ideogram4 import PRESETS, Ideogram4Pipeline, Ideogram4PipelineConfig + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run a local-caption Ideogram 4 split-device smoke test." + ) + parser.add_argument( + "--caption-file", + type=Path, + required=True, + help="UTF-8 file containing an already-structured Ideogram JSON caption.", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--weights-repo", default="ideogram-ai/ideogram-4-nf4") + parser.add_argument("--diffusion-device", default="cuda:0") + parser.add_argument("--text-device", default="cuda:1") + parser.add_argument("--height", type=int, default=256) + parser.add_argument("--width", type=int, default=256) + parser.add_argument( + "--sampler-preset", + choices=["custom", *sorted(PRESETS)], + default="custom", + help="Use an official preset, or custom for --num-steps/--guidance-scale.", + ) + parser.add_argument("--num-steps", type=int, default=2) + parser.add_argument("--guidance-scale", type=float, default=7.0) + parser.add_argument("--seed", type=int, default=0) + return parser.parse_args() + + +def _print_preflight(diffusion_device: str, text_device: str) -> None: + print(f"torch={torch.__version__}") + print(f"HIP={torch.version.hip}") + print( + "allocator_config=" + f"{os.environ.get('PYTORCH_ALLOC_CONF') or os.environ.get('PYTORCH_CUDA_ALLOC_CONF') or ''}" + ) + try: + print(f"allocator_backend={torch.cuda.get_allocator_backend()}") + except Exception as exc: # noqa: BLE001 - diagnostics must not block validation + print(f"allocator_backend=") + print(f"MIOPEN_FIND_MODE={os.environ.get('MIOPEN_FIND_MODE', '')}") + print( + f"MIOPEN_DEBUG_CONV_GEMM={os.environ.get('MIOPEN_DEBUG_CONV_GEMM', '')}" + ) + print(f"MIOPEN_DEBUG_CONV_FFT={os.environ.get('MIOPEN_DEBUG_CONV_FFT', '')}") + print(f"visible_device_count={torch.cuda.device_count()}") + for index in range(torch.cuda.device_count()): + properties = torch.cuda.get_device_properties(index) + print( + f"cuda:{index}: {torch.cuda.get_device_name(index)}, " + f"VRAM={properties.total_memory / 1024**3:.2f} GiB" + ) + + resolved_diffusion = torch.device(diffusion_device) + resolved_text = torch.device(text_device) + if resolved_diffusion != resolved_text: + if resolved_diffusion.index is None or resolved_text.index is None: + raise ValueError("Split-device smoke tests require explicit CUDA indexes") + required_count = max(resolved_diffusion.index, resolved_text.index) + 1 + if torch.cuda.device_count() < required_count: + raise RuntimeError( + f"Requested {resolved_diffusion} and {resolved_text}, but only " + f"{torch.cuda.device_count()} CUDA/ROCm device(s) are visible" + ) + print( + f"peer_access {resolved_diffusion}->{resolved_text}=" + f"{torch.cuda.can_device_access_peer(resolved_diffusion.index, resolved_text.index)}" + ) + print( + f"peer_access {resolved_text}->{resolved_diffusion}=" + f"{torch.cuda.can_device_access_peer(resolved_text.index, resolved_diffusion.index)}" + ) + + +def main() -> None: + args = _parse_args() + _print_preflight(args.diffusion_device, args.text_device) + caption = args.caption_file.read_text(encoding="utf-8") + + pipeline = Ideogram4Pipeline.from_pretrained( + config=Ideogram4PipelineConfig(weights_repo=args.weights_repo), + device=args.diffusion_device, + text_device=args.text_device, + dtype=torch.bfloat16, + ) + + tracked_devices = { + torch.device(args.diffusion_device), + torch.device(args.text_device), + } + for device in tracked_devices: + torch.cuda.reset_peak_memory_stats(device) + + generation_kwargs = { + "num_steps": args.num_steps, + "guidance_scale": args.guidance_scale, + } + if args.sampler_preset != "custom": + preset = PRESETS[args.sampler_preset] + generation_kwargs = { + "num_steps": preset.num_steps, + "guidance_schedule": preset.guidance_schedule, + "mu": preset.mu, + "std": preset.std, + } + print(f"generation={args.width}x{args.height}, sampler={args.sampler_preset}") + images = pipeline( + caption, + height=args.height, + width=args.width, + seed=args.seed, + **generation_kwargs, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + images[0].save(args.output) + print(f"saved={args.output.resolve()}") + + for device in sorted(tracked_devices, key=str): + peak_allocated = torch.cuda.max_memory_allocated(device) / 1024**3 + peak_reserved = torch.cuda.max_memory_reserved(device) / 1024**3 + print( + f"peak_memory device={device}, name={torch.cuda.get_device_name(device)}, " + f"allocated={peak_allocated:.3f} GiB, reserved={peak_reserved:.3f} GiB" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_comfyui_api.py b/scripts/validate_comfyui_api.py new file mode 100644 index 0000000..b7b3494 --- /dev/null +++ b/scripts/validate_comfyui_api.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import argparse +import json +import time +import urllib.error +import urllib.parse +import urllib.request +import uuid +from pathlib import Path +from typing import Any + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Submit and verify the Ideogram 4 ComfyUI acceptance workflow." + ) + parser.add_argument("--base-url", default="http://127.0.0.1:8189") + parser.add_argument("--caption-file", type=Path, required=True) + parser.add_argument("--startup-timeout", type=float, default=300.0) + parser.add_argument("--generation-timeout", type=float, default=3600.0) + return parser.parse_args() + + +def _request_json( + url: str, + *, + payload: dict[str, Any] | None = None, + timeout: float = 30.0, +) -> dict[str, Any]: + data = None + headers = {} + if payload is not None: + data = json.dumps(payload).encode("utf-8") + headers["Content-Type"] = "application/json" + request = urllib.request.Request(url, data=data, headers=headers) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + body = response.read().decode("utf-8") + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + raise RuntimeError(f"ComfyUI HTTP {exc.code} from {url}: {body}") from exc + result = json.loads(body) + if not isinstance(result, dict): + raise TypeError(f"Expected a JSON object from {url}, got {type(result).__name__}") + return result + + +def _wait_for_server(base_url: str, timeout: float) -> None: + deadline = time.monotonic() + timeout + last_error: Exception | None = None + while time.monotonic() < deadline: + try: + _request_json(f"{base_url}/system_stats", timeout=5.0) + print(f"ComfyUI API ready: {base_url}", flush=True) + return + except (OSError, RuntimeError, ValueError) as exc: + last_error = exc + time.sleep(1.0) + raise TimeoutError(f"ComfyUI did not become ready within {timeout:g}s: {last_error}") + + +def _workflow(caption: str) -> dict[str, Any]: + return { + "1": { + "class_type": "Ideogram4PipelineLoader", + "inputs": {"model_weights": "4.0 NF4"}, + }, + "2": { + "class_type": "Ideogram4Generate", + "inputs": { + "pipeline": ["1", 0], + "prompt": caption, + "width": 1024, + "height": 1024, + "sampler_preset": "4.0 Default 20", + "num_steps": 20, + "guidance_scale": 7.0, + "mu": 0.0, + "std": 1.75, + "seed": 0, + }, + }, + "3": { + "class_type": "SaveImage", + "inputs": { + "images": ["2", 0], + "filename_prefix": "ideogram4-split-validation", + }, + }, + } + + +def _error_details(entry: dict[str, Any]) -> str: + status = entry.get("status", {}) + messages = status.get("messages", []) if isinstance(status, dict) else [] + for message in reversed(messages): + if ( + isinstance(message, list) + and len(message) >= 2 + and message[0] == "execution_error" + ): + return json.dumps(message[1], indent=2) + return json.dumps(status, indent=2) + + +def _wait_for_result(base_url: str, prompt_id: str, timeout: float) -> dict[str, Any]: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + history = _request_json(f"{base_url}/history/{prompt_id}") + entry = history.get(prompt_id) + if isinstance(entry, dict): + status = entry.get("status", {}) + if isinstance(status, dict) and status.get("completed"): + if status.get("status_str") != "success": + raise RuntimeError("ComfyUI workflow failed:\n" + _error_details(entry)) + return entry + time.sleep(2.0) + raise TimeoutError(f"ComfyUI workflow {prompt_id} exceeded {timeout:g}s") + + +def _saved_images(entry: dict[str, Any]) -> list[dict[str, Any]]: + outputs = entry.get("outputs", {}) + save_output = outputs.get("3", {}) if isinstance(outputs, dict) else {} + images = save_output.get("images", []) if isinstance(save_output, dict) else [] + return [image for image in images if isinstance(image, dict)] + + +def _verify_image(base_url: str, image: dict[str, Any]) -> int: + query = urllib.parse.urlencode( + { + "filename": image.get("filename", ""), + "subfolder": image.get("subfolder", ""), + "type": image.get("type", "output"), + } + ) + with urllib.request.urlopen(f"{base_url}/view?{query}", timeout=30.0) as response: + data = response.read() + if not data: + raise RuntimeError(f"ComfyUI returned an empty saved image: {image}") + return len(data) + + +def main() -> None: + args = _parse_args() + base_url = args.base_url.rstrip("/") + caption = args.caption_file.read_text(encoding="utf-8") + json.loads(caption) + + _wait_for_server(base_url, args.startup_timeout) + queued = _request_json( + f"{base_url}/prompt", + payload={"prompt": _workflow(caption), "client_id": str(uuid.uuid4())}, + ) + prompt_id = queued.get("prompt_id") + if not isinstance(prompt_id, str) or not prompt_id: + raise RuntimeError(f"ComfyUI did not return a prompt_id: {queued}") + if queued.get("node_errors"): + raise RuntimeError( + "ComfyUI rejected workflow nodes: " + json.dumps(queued["node_errors"], indent=2) + ) + print(f"ComfyUI workflow queued: {prompt_id}", flush=True) + + entry = _wait_for_result(base_url, prompt_id, args.generation_timeout) + images = _saved_images(entry) + if not images: + raise RuntimeError("ComfyUI workflow completed without a SaveImage output") + size = _verify_image(base_url, images[0]) + print(f"ComfyUI saved image: {json.dumps(images[0], sort_keys=True)}", flush=True) + print(f"ComfyUI saved image bytes: {size}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_rocm_split_devices.sh b/scripts/validate_rocm_split_devices.sh new file mode 100755 index 0000000..e75e5e7 --- /dev/null +++ b/scripts/validate_rocm_split_devices.sh @@ -0,0 +1,469 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +usage() { + cat <<'EOF' +Validate Ideogram 4 split-device inference on the target Fedora/ROCm host. + +Usage: + scripts/validate_rocm_split_devices.sh --caption FILE [options] + +Required: + --caption FILE Structured Ideogram JSON caption stored locally. + +Options: + --python PATH Python from the ROCm/ComfyUI environment. + --comfyui-dir DIR ComfyUI checkout; also selects its .venv by default. + --output-dir DIR Logs and images directory (default: timestamped). + --diffusion-device DEVICE Diffusion/VAE device (default: cuda:0). + --text-device DEVICE Qwen3-VL device (default: cuda:1). + --visible-devices LIST ROCR_VISIBLE_DEVICES value (default: existing or 0,1). + --weights-repo REPO Hugging Face weights repo. + --expected-diffusion TEXT Required substring in diffusion device name. + Default: 7900 XT + --expected-text TEXT Required substring in text device name. + Default: 8060S + --allow-name-mismatch Warn instead of failing when names do not match. + --skip-small Skip the 256x256 two-step smoke test. + --skip-1024 Skip the 1024x1024 V4_DEFAULT_20 test. + --launch-comfyui Launch ComfyUI after core tests (foreground). + --comfyui-port PORT ComfyUI port (default: 8189). + -h, --help Show this help. + +The script never calls a hosted Magic Prompt API. It installs this core checkout +editable with --no-deps so it cannot replace the existing ROCm PyTorch build. +EOF +} + +fail() { + echo "ERROR: $*" >&2 + exit 1 +} + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +CORE_DIR="$(cd -- "${SCRIPT_DIR}/.." && pwd)" +[[ -f "$CORE_DIR/src/ideogram4/pipeline_ideogram4.py" ]] || fail \ + "run this script from the modified ideogram4 checkout" +CAPTION_FILE="" +PYTHON_BIN="" +COMFYUI_DIR="" +OUTPUT_DIR="" +DIFFUSION_DEVICE="cuda:0" +TEXT_DEVICE="cuda:1" +VISIBLE_DEVICES="${ROCR_VISIBLE_DEVICES:-0,1}" +WEIGHTS_REPO="ideogram-ai/ideogram-4-nf4" +EXPECTED_DIFFUSION="7900 XT" +EXPECTED_TEXT="8060S" +ALLOW_NAME_MISMATCH=0 +SKIP_SMALL=0 +SKIP_1024=0 +LAUNCH_COMFYUI=0 +COMFYUI_PORT=8189 + +while (($#)); do + case "$1" in + --caption) + (($# >= 2)) || fail "$1 requires a value" + CAPTION_FILE="$2" + shift 2 + ;; + --python) + (($# >= 2)) || fail "$1 requires a value" + PYTHON_BIN="$2" + shift 2 + ;; + --comfyui-dir) + (($# >= 2)) || fail "$1 requires a value" + COMFYUI_DIR="$2" + shift 2 + ;; + --output-dir) + (($# >= 2)) || fail "$1 requires a value" + OUTPUT_DIR="$2" + shift 2 + ;; + --diffusion-device) + (($# >= 2)) || fail "$1 requires a value" + DIFFUSION_DEVICE="$2" + shift 2 + ;; + --text-device) + (($# >= 2)) || fail "$1 requires a value" + TEXT_DEVICE="$2" + shift 2 + ;; + --visible-devices) + (($# >= 2)) || fail "$1 requires a value" + VISIBLE_DEVICES="$2" + shift 2 + ;; + --weights-repo) + (($# >= 2)) || fail "$1 requires a value" + WEIGHTS_REPO="$2" + shift 2 + ;; + --expected-diffusion) + (($# >= 2)) || fail "$1 requires a value" + EXPECTED_DIFFUSION="$2" + shift 2 + ;; + --expected-text) + (($# >= 2)) || fail "$1 requires a value" + EXPECTED_TEXT="$2" + shift 2 + ;; + --allow-name-mismatch) + ALLOW_NAME_MISMATCH=1 + shift + ;; + --skip-small) + SKIP_SMALL=1 + shift + ;; + --skip-1024) + SKIP_1024=1 + shift + ;; + --launch-comfyui) + LAUNCH_COMFYUI=1 + shift + ;; + --comfyui-port) + (($# >= 2)) || fail "$1 requires a value" + COMFYUI_PORT="$2" + shift 2 + ;; + -h|--help) + usage + exit 0 + ;; + *) + fail "unknown argument: $1" + ;; + esac +done + +[[ "$(uname -s)" == "Linux" ]] || fail "this script must run on Linux" +[[ -n "$CAPTION_FILE" ]] || fail "--caption is required" +[[ -r "$CAPTION_FILE" ]] || fail "caption file is not readable: $CAPTION_FILE" +CAPTION_FILE="$(readlink -f -- "$CAPTION_FILE")" + +if [[ -n "$COMFYUI_DIR" ]]; then + COMFYUI_DIR="$(readlink -f -- "$COMFYUI_DIR")" + [[ -f "$COMFYUI_DIR/main.py" ]] || fail "ComfyUI main.py not found under $COMFYUI_DIR" +fi + +if [[ -z "$PYTHON_BIN" && -n "$COMFYUI_DIR" ]]; then + if [[ -x "$COMFYUI_DIR/.venv/bin/python" ]]; then + PYTHON_BIN="$COMFYUI_DIR/.venv/bin/python" + elif [[ -x "$COMFYUI_DIR/venv/bin/python" ]]; then + PYTHON_BIN="$COMFYUI_DIR/venv/bin/python" + fi +fi +PYTHON_BIN="${PYTHON_BIN:-python3}" +PYTHON_BIN="$(command -v -- "$PYTHON_BIN" || true)" +[[ -n "$PYTHON_BIN" && -x "$PYTHON_BIN" ]] || fail "Python executable not found" + +if [[ -z "$OUTPUT_DIR" ]]; then + OUTPUT_DIR="$HOME/ideogram4-rocm-validation-$(date +%Y%m%d-%H%M%S)" +fi +mkdir -p -- "$OUTPUT_DIR" +OUTPUT_DIR="$(readlink -f -- "$OUTPUT_DIR")" + +export ROCR_VISIBLE_DEVICES="$VISIBLE_DEVICES" +export IDEOGRAM4_DIFFUSION_DEVICE="$DIFFUSION_DEVICE" +export IDEOGRAM4_TEXT_DEVICE="$TEXT_DEVICE" +export IDEOGRAM4_REPO="$CORE_DIR" +export PYTHONUNBUFFERED=1 + +# ROCm 7.2 does not support PyTorch's expandable allocator segments. Keep +# MIOpen away from its high-workspace GEMM/FFT convolution paths instead; the +# direct, Winograd, and implicit-GEMM solvers remain available for VAE decode. +export MIOPEN_FIND_MODE="${MIOPEN_FIND_MODE:-FAST}" +export MIOPEN_DEBUG_CONV_GEMM="${MIOPEN_DEBUG_CONV_GEMM:-0}" +export MIOPEN_DEBUG_CONV_FFT="${MIOPEN_DEBUG_CONV_FFT:-0}" + +SUMMARY_LOG="$OUTPUT_DIR/summary.log" +SMALL_LOG="$OUTPUT_DIR/core-256.log" +FULL_LOG="$OUTPUT_DIR/core-1024.log" +GPU_LOG="$OUTPUT_DIR/gpu-monitor.log" +MONITOR_PID="" +COMFYUI_PID="" + +stop_monitor() { + if [[ -n "$MONITOR_PID" ]] && kill -0 "$MONITOR_PID" 2>/dev/null; then + kill "$MONITOR_PID" 2>/dev/null || true + wait "$MONITOR_PID" 2>/dev/null || true + fi +} + +stop_comfyui() { + if [[ -n "$COMFYUI_PID" ]] && kill -0 "$COMFYUI_PID" 2>/dev/null; then + kill "$COMFYUI_PID" 2>/dev/null || true + wait "$COMFYUI_PID" 2>/dev/null || true + fi +} + +on_exit() { + status=$? + stop_monitor + stop_comfyui + if ((status == 0)); then + echo "Validation command completed. Artifacts: $OUTPUT_DIR" + else + echo "Validation failed with exit code $status. Logs: $OUTPUT_DIR" >&2 + fi +} +trap on_exit EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +start_monitor() { + if ! command -v rocm-smi >/dev/null 2>&1; then + echo "rocm-smi not found; relying on PyTorch peak-memory reporting" | tee -a "$SUMMARY_LOG" + return + fi + ( + while true; do + date --iso-8601=seconds + rocm-smi --showproductname --showuse --showmemuse || true + sleep 1 + done + ) >"$GPU_LOG" 2>&1 & + MONITOR_PID=$! + echo "GPU monitor PID=$MONITOR_PID log=$GPU_LOG" | tee -a "$SUMMARY_LOG" +} + +{ + echo "Ideogram 4 ROCm split-device validation" + echo "started=$(date --iso-8601=seconds)" + echo "host=$(hostname)" + echo "kernel=$(uname -srmo)" + echo "core_dir=$CORE_DIR" + echo "python=$PYTHON_BIN" + echo "caption=$CAPTION_FILE" + echo "output_dir=$OUTPUT_DIR" + echo "ROCR_VISIBLE_DEVICES=$ROCR_VISIBLE_DEVICES" + echo "diffusion_device=$DIFFUSION_DEVICE" + echo "text_device=$TEXT_DEVICE" + echo "weights_repo=$WEIGHTS_REPO" + echo "HF_HOME=${HF_HOME:-}" + echo "PYTORCH_ALLOC_CONF=${PYTORCH_ALLOC_CONF:-}" + echo "PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-}" + echo "MIOPEN_FIND_MODE=$MIOPEN_FIND_MODE" + echo "MIOPEN_DEBUG_CONV_GEMM=$MIOPEN_DEBUG_CONV_GEMM" + echo "MIOPEN_DEBUG_CONV_FFT=$MIOPEN_DEBUG_CONV_FFT" + df -h -- "$OUTPUT_DIR" +} | tee "$SUMMARY_LOG" + +"$PYTHON_BIN" - "$DIFFUSION_DEVICE" "$TEXT_DEVICE" \ + "$EXPECTED_DIFFUSION" "$EXPECTED_TEXT" "$ALLOW_NAME_MISMATCH" <<'PY' | tee -a "$SUMMARY_LOG" +import sys + +try: + import torch +except ImportError as exc: + raise SystemExit( + "PyTorch is missing from the selected environment. Install the ROCm build first." + ) from exc + +diffusion = torch.device(sys.argv[1]) +text = torch.device(sys.argv[2]) +expected_diffusion = sys.argv[3].casefold() +expected_text = sys.argv[4].casefold() +allow_mismatch = bool(int(sys.argv[5])) + +print("torch:", torch.__version__) +print("HIP:", torch.version.hip) +print("CUDA/ROCm available:", torch.cuda.is_available()) +print("device count:", torch.cuda.device_count()) + +if torch.version.hip is None: + raise SystemExit("Selected Python does not contain a ROCm PyTorch build (torch.version.hip is None)") +if not torch.cuda.is_available(): + raise SystemExit("ROCm devices are not available through torch.cuda") + +for index in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(index) + print( + f"cuda:{index}: {torch.cuda.get_device_name(index)} " + f"VRAM={props.total_memory / 1024**3:.2f} GiB" + ) + +if diffusion.type != "cuda" or text.type != "cuda": + raise SystemExit(f"Expected CUDA/ROCm roles, got {diffusion} and {text}") +if diffusion.index is None or text.index is None: + raise SystemExit("Split-device validation requires explicit cuda:N indexes") +highest = max(diffusion.index, text.index) +if torch.cuda.device_count() <= highest: + raise SystemExit( + f"Requested {diffusion} and {text}, but only {torch.cuda.device_count()} device(s) are visible" + ) +if diffusion == text: + raise SystemExit("Diffusion and text devices must differ for this validation") + +diffusion_name = torch.cuda.get_device_name(diffusion.index) +text_name = torch.cuda.get_device_name(text.index) +diffusion_vram_gib = torch.cuda.get_device_properties(diffusion.index).total_memory / 1024**3 +text_vram_gib = torch.cuda.get_device_properties(text.index).total_memory / 1024**3 +mismatches = [] +if expected_diffusion and expected_diffusion not in diffusion_name.casefold(): + generic_amd_name = diffusion_name.casefold() == "amd radeon graphics" + expected_7900_memory = 18.0 <= diffusion_vram_gib <= 24.0 + if expected_diffusion == "7900 xt" and generic_amd_name and expected_7900_memory: + print( + "WARNING: ROCm reported the diffusion GPU with the generic name " + f"{diffusion_name!r}; accepting it as the expected 7900 XT based on " + f"its {diffusion_vram_gib:.2f} GiB VRAM." + ) + else: + mismatches.append( + f"diffusion device {diffusion} is {diffusion_name!r} " + f"with {diffusion_vram_gib:.2f} GiB VRAM, expected substring {sys.argv[3]!r}" + ) +if expected_text and expected_text not in text_name.casefold(): + mismatches.append( + f"text device {text} is {text_name!r} with {text_vram_gib:.2f} GiB VRAM, " + f"expected substring {sys.argv[4]!r}" + ) +if mismatches: + message = "\n".join(mismatches) + if allow_mismatch: + print("WARNING: device-name mismatch allowed:\n" + message) + else: + raise SystemExit( + "Device mapping does not match the requested roles:\n" + + message + + "\nCorrect the device arguments or pass --allow-name-mismatch intentionally." + ) + +print(f"resolved diffusion: {diffusion} ({diffusion_name})") +print(f"resolved text: {text} ({text_name})") +print( + f"peer access {diffusion}->{text}:", + torch.cuda.can_device_access_peer(diffusion.index, text.index), +) +print( + f"peer access {text}->{diffusion}:", + torch.cuda.can_device_access_peer(text.index, diffusion.index), +) +PY + +echo "Installing editable core without changing environment dependencies..." | tee -a "$SUMMARY_LOG" +"$PYTHON_BIN" -m pip install --no-deps -e "$CORE_DIR" 2>&1 | tee -a "$SUMMARY_LOG" + +"$PYTHON_BIN" - <<'PY' | tee -a "$SUMMARY_LOG" +import importlib +import ideogram4 + +required = ( + "accelerate", + "bitsandbytes", + "einops", + "huggingface_hub", + "PIL", + "safetensors", + "sentencepiece", + "transformers", +) +missing = [] +for name in required: + try: + importlib.import_module(name) + except ImportError: + missing.append(name) +if missing: + raise SystemExit("Missing dependencies in selected environment: " + ", ".join(missing)) + +print("ideogram4 import:", ideogram4.__file__) +try: + from huggingface_hub import get_token + print("Hugging Face authentication available:", bool(get_token())) +except Exception as exc: + print("WARNING: unable to inspect Hugging Face authentication:", exc) +PY + +start_monitor + +if ((SKIP_SMALL == 0)); then + echo "Running 256x256 two-step core smoke test..." | tee -a "$SUMMARY_LOG" + "$PYTHON_BIN" "$CORE_DIR/scripts/smoke_split_devices.py" \ + --caption-file "$CAPTION_FILE" \ + --output "$OUTPUT_DIR/ideogram4-256.png" \ + --weights-repo "$WEIGHTS_REPO" \ + --diffusion-device "$DIFFUSION_DEVICE" \ + --text-device "$TEXT_DEVICE" \ + --height 256 --width 256 --num-steps 2 \ + 2>&1 | tee "$SMALL_LOG" +fi + +if ((SKIP_1024 == 0)); then + echo "Running 1024x1024 V4_DEFAULT_20 acceptance test..." | tee -a "$SUMMARY_LOG" + "$PYTHON_BIN" "$CORE_DIR/scripts/smoke_split_devices.py" \ + --caption-file "$CAPTION_FILE" \ + --output "$OUTPUT_DIR/ideogram4-1024.png" \ + --weights-repo "$WEIGHTS_REPO" \ + --diffusion-device "$DIFFUSION_DEVICE" \ + --text-device "$TEXT_DEVICE" \ + --height 1024 --width 1024 \ + --sampler-preset V4_DEFAULT_20 \ + 2>&1 | tee "$FULL_LOG" +fi + +if grep -Fq "staging through CPU" "$SMALL_LOG" "$FULL_LOG" 2>/dev/null; then + COPY_PATH="CPU-staged fallback" +else + COPY_PATH="direct device transfer (no CPU-staging warning observed)" +fi +echo "feature_copy_path=$COPY_PATH" | tee -a "$SUMMARY_LOG" +echo "core_validation=PASS" | tee -a "$SUMMARY_LOG" + +if ((LAUNCH_COMFYUI == 1)); then + [[ -n "$COMFYUI_DIR" ]] || fail "--launch-comfyui requires --comfyui-dir" + WRAPPER_DIR="$COMFYUI_DIR/custom_nodes/ComfyUI-Ideogram4" + if [[ ! -f "$WRAPPER_DIR/nodes.py" ]]; then + fail "ComfyUI-Ideogram4 is not installed under $COMFYUI_DIR/custom_nodes" + fi + if ! grep -Fq "IDEOGRAM4_TEXT_DEVICE" "$WRAPPER_DIR/nodes.py"; then + fail "installed ComfyUI-Ideogram4 does not contain the split-device wrapper change" + fi + echo "Launching ComfyUI on port $COMFYUI_PORT." | tee -a "$SUMMARY_LOG" + echo "Submitting 4.0 NF4, 1024x1024, V4_DEFAULT_20 through the ComfyUI API." \ + | tee -a "$SUMMARY_LOG" + cd -- "$COMFYUI_DIR" + "$PYTHON_BIN" main.py \ + --listen 0.0.0.0 \ + --port "$COMFYUI_PORT" \ + --disable-pinned-memory \ + >"$OUTPUT_DIR/comfyui.log" 2>&1 & + COMFYUI_PID=$! + cd -- "$CORE_DIR" + if ! "$PYTHON_BIN" "$CORE_DIR/scripts/validate_comfyui_api.py" \ + --base-url "http://127.0.0.1:$COMFYUI_PORT" \ + --caption-file "$CAPTION_FILE" \ + 2>&1 | tee "$OUTPUT_DIR/comfyui-api-validation.log"; then + echo "ComfyUI API validation failed. Recent server log:" >&2 + tail -n 200 "$OUTPUT_DIR/comfyui.log" >&2 || true + exit 1 + fi + echo "comfyui_validation=PASS" | tee -a "$SUMMARY_LOG" + echo "ComfyUI remains available at http://0.0.0.0:$COMFYUI_PORT; press Ctrl-C to stop it." \ + | tee -a "$SUMMARY_LOG" + wait "$COMFYUI_PID" +else + if [[ -n "$COMFYUI_DIR" ]]; then + cat < str: + """Return a useful device name without allowing diagnostics to fail.""" + try: + if device.type == "cuda" and device.index is not None: + return torch.cuda.get_device_name(device.index) + except Exception: # noqa: BLE001, S110 - diagnostics must never fail inference + pass + return str(device) + + +def _device_description(device: torch.device) -> str: + return f"{device} ({_device_name(device)})" + + +def _synchronize_device(device: torch.device) -> None: + try: + if device.type == "cuda": + torch.cuda.synchronize(device) + except Exception: # noqa: BLE001, S110 - timing must never fail inference + pass + + +def _log_device_memory(devices: Sequence[torch.device]) -> None: + seen: set[str] = set() + for device in devices: + key = str(device) + if key in seen: + continue + seen.add(key) + try: + if device.type != "cuda": + continue + allocated = torch.cuda.memory_allocated(device) / 1024**2 + reserved = torch.cuda.memory_reserved(device) / 1024**2 + print( + f"Ideogram4 memory: device={_device_description(device)}, " + f"allocated={allocated:.1f} MiB, reserved={reserved:.1f} MiB", + flush=True, + ) + except Exception as exc: # noqa: BLE001 - diagnostics must never fail inference + warnings.warn( + f"Unable to read Ideogram4 memory diagnostics for {device}: {exc}", + stacklevel=2, + ) + + +def _validate_device_configuration( + diffusion_device: torch.device, text_device: torch.device +) -> None: + if diffusion_device == text_device: + return + if diffusion_device.type == "cuda" and text_device.type == "cuda": + if diffusion_device.index is None or text_device.index is None: + raise ValueError( + "Separate Ideogram4 CUDA/ROCm devices must use explicit indexes; got " + f"diffusion_device={diffusion_device}, text_device={text_device}" + ) + try: + device_count = torch.cuda.device_count() + except Exception as exc: + raise RuntimeError( + "Unable to inspect CUDA/ROCm devices for split Ideogram4 placement: " + f"diffusion_device={diffusion_device}, text_device={text_device}" + ) from exc + highest_index = max(diffusion_device.index, text_device.index) + if device_count <= highest_index: + raise RuntimeError( + "Split Ideogram4 placement requested unavailable CUDA/ROCm devices: " + f"diffusion_device={diffusion_device}, text_device={text_device}, " + f"visible_device_count={device_count}" + ) + + +def _move_tensor_to_device( + tensor: torch.Tensor, destination: torch.device +) -> torch.Tensor: + """Move one tensor directly, with a CPU-staged fallback for GPU peers.""" + if tensor.device == destination: + return tensor + try: + return tensor.to(destination) + except RuntimeError as direct_error: + warnings.warn( + f"Direct tensor transfer {tensor.device} -> {destination} failed; " + f"staging through CPU. {direct_error}", + stacklevel=2, + ) + return tensor.to("cpu").to(destination) + + +def _append_image_feature_zeros( + text_features: torch.Tensor, + num_image_tokens: int, + destination: torch.device, +) -> torch.Tensor: + """Rebuild full diffusion conditioning after compact text encoding.""" + batch_size, _, feature_dim = text_features.shape + image_feature_zeros = torch.zeros( + batch_size, + num_image_tokens, + feature_dim, + dtype=text_features.dtype, + device=destination, + ) + return torch.cat([text_features, image_feature_zeros], dim=1) + + def _load_subfolder_state_dict( repo_id: str, subfolder: str, basename: str ) -> dict[str, torch.Tensor]: @@ -259,6 +369,7 @@ def __init__( config: Ideogram4PipelineConfig, device: torch.device, dtype: torch.dtype, + text_device: torch.device | None = None, ) -> None: self.conditional_transformer = conditional_transformer self.unconditional_transformer = unconditional_transformer @@ -266,13 +377,17 @@ def __init__( self.text_tokenizer = text_tokenizer self.autoencoder = autoencoder self.config = config - self.device = device + self.diffusion_device = torch.device(device) + self.text_device = ( + torch.device(text_device) if text_device is not None else self.diffusion_device + ) + self.device = self.diffusion_device self.dtype = dtype self.caption_verifier = CaptionVerifier() shift, scale = get_latent_norm() - self.latent_shift = shift.to(device) - self.latent_scale = scale.to(device) + self.latent_shift = shift.to(self.diffusion_device) + self.latent_scale = scale.to(self.diffusion_device) @classmethod def from_pretrained( @@ -280,12 +395,23 @@ def from_pretrained( *, config: Optional[Ideogram4PipelineConfig] = None, device: str | torch.device = "cuda", + text_device: str | torch.device | None = None, dtype: torch.dtype = torch.bfloat16, transformer_config: Optional[Ideogram4Config] = None, ) -> "Ideogram4Pipeline": config = config or Ideogram4PipelineConfig() transformer_config = transformer_config or Ideogram4Config() - device = torch.device(device) + diffusion_device = torch.device(device) + resolved_text_device = ( + torch.device(text_device) if text_device is not None else diffusion_device + ) + _validate_device_configuration(diffusion_device, resolved_text_device) + print( + "Ideogram4 devices: " + f"diffusion_device={_device_description(diffusion_device)}, " + f"text_device={_device_description(resolved_text_device)}", + flush=True, + ) conditional_state_dict = _load_indexed_or_single_state_dict( config.weights_repo, config.conditional_index_filename @@ -298,33 +424,36 @@ def from_pretrained( ) conditional_transformer = _build_transformer( - transformer_config, conditional_state_dict, device, dtype + transformer_config, conditional_state_dict, diffusion_device, dtype ) del conditional_state_dict unconditional_transformer = _build_transformer( - transformer_config, unconditional_state_dict, device, dtype + transformer_config, unconditional_state_dict, diffusion_device, dtype ) del unconditional_state_dict text_tokenizer, text_encoder = _load_qwen3_vl( config.weights_repo, - device, + resolved_text_device, dtype, tokenizer_subfolder=config.tokenizer_subfolder, text_encoder_subfolder=config.text_encoder_subfolder, ) - autoencoder = _load_autoencoder(autoencoder_weights, device, dtype) + autoencoder = _load_autoencoder(autoencoder_weights, diffusion_device, dtype) - return cls( + pipeline = cls( conditional_transformer=conditional_transformer, unconditional_transformer=unconditional_transformer, text_encoder=text_encoder, text_tokenizer=text_tokenizer, autoencoder=autoencoder, config=config, - device=device, + device=diffusion_device, dtype=dtype, + text_device=resolved_text_device, ) + _log_device_memory([diffusion_device, resolved_text_device]) + return pipeline def _tokenize(self, prompt: str) -> tuple[torch.Tensor, int]: """Build chat-formatted token ids for a single prompt.""" @@ -400,11 +529,12 @@ def _build_inputs( segment_ids[b, offset : offset + total_unpadded] = 1 return { - "token_ids": token_ids.to(self.device), - "text_position_ids": text_position_ids.to(self.device), - "position_ids": position_ids.to(self.device), - "segment_ids": segment_ids.to(self.device), - "indicator": indicator.to(self.device), + "token_ids": token_ids.to(self.text_device), + "text_position_ids": text_position_ids.to(self.text_device), + "text_indicator": indicator.to(self.text_device), + "position_ids": position_ids.to(self.diffusion_device), + "segment_ids": segment_ids.to(self.diffusion_device), + "indicator": indicator.to(self.diffusion_device), "num_image_tokens": num_image_tokens, # type: ignore[dict-item] "grid_h": grid_h, # type: ignore[dict-item] "grid_w": grid_w, # type: ignore[dict-item] @@ -454,19 +584,31 @@ def _encode_text( self, token_ids: torch.Tensor, text_position_ids: torch.Tensor, - indicator: torch.Tensor, + text_indicator: torch.Tensor, + max_text_tokens: int, ) -> torch.Tensor: """Run Qwen3-VL and stack hidden states from the activation layers. - Returns a (B, L, hidden_size * num_layers) float32 tensor. + Returns text-position features on the diffusion device as a + (B, max_text_tokens, hidden_size * num_layers) float32 tensor. """ + token_ids = token_ids[:, :max_text_tokens].contiguous() + text_position_ids = text_position_ids[:, :max_text_tokens].contiguous() + text_indicator = text_indicator[:, :max_text_tokens].contiguous() batch_size, seq_len = token_ids.shape # Real text positions are exactly the LLM_TOKEN_INDICATOR positions. - attention_mask = (indicator == LLM_TOKEN_INDICATOR).to(torch.long) + attention_mask = (text_indicator == LLM_TOKEN_INDICATOR).to(torch.long) pos_2d = text_position_ids[..., 0].contiguous() + print( + "Ideogram4 text encoding started: " + f"sequence_length={seq_len}, device={_device_description(self.text_device)}", + flush=True, + ) + _synchronize_device(self.text_device) + encode_started = time.perf_counter() with torch.no_grad(): selected = self._get_qwen3_vl_embeddings(token_ids, attention_mask, pos_2d) stacked = torch.stack(selected, dim=0) # (num_taps, B, L, H) @@ -477,7 +619,24 @@ def _encode_text( # text features at LLM_TOKEN_INDICATOR positions. text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1) stacked = stacked * text_mask - return stacked.to(torch.float32) + stacked = stacked.to(torch.float32) + _synchronize_device(self.text_device) + encode_elapsed = time.perf_counter() - encode_started + + transfer_mib = stacked.numel() * stacked.element_size() / 1024**2 + copy_started = time.perf_counter() + text_features = _move_tensor_to_device(stacked, self.diffusion_device) + _synchronize_device(self.diffusion_device) + copy_elapsed = time.perf_counter() - copy_started + print( + "Ideogram4 text encoding finished: " + f"feature_shape={tuple(text_features.shape)}, dtype={text_features.dtype}, " + f"transfer={transfer_mib:.1f} MiB, source={self.text_device}, " + f"destination={self.diffusion_device}, encode_time={encode_elapsed:.3f}s, " + f"copy_time={copy_elapsed:.3f}s", + flush=True, + ) + return text_features def _verify_prompts( self, prompts: list[str], *, raise_on_issues: bool = True @@ -525,11 +684,11 @@ def __call__( schedule = schedule or get_schedule_for_resolution( (height, width), known_mean=mu, std=std ) - step_intervals = make_step_intervals(num_steps).to(self.device) + step_intervals = make_step_intervals(num_steps).to(self.diffusion_device) if guidance_schedule is not None: gw_per_step = torch.as_tensor( - guidance_schedule, dtype=torch.float32, device=self.device + guidance_schedule, dtype=torch.float32, device=self.diffusion_device ) if gw_per_step.shape != (num_steps,): raise ValueError( @@ -538,7 +697,10 @@ def __call__( ) else: gw_per_step = torch.full( - (num_steps,), float(guidance_scale), dtype=torch.float32, device=self.device + (num_steps,), + float(guidance_scale), + dtype=torch.float32, + device=self.diffusion_device, ) inputs = self._build_inputs(prompts, height=height, width=width) @@ -549,7 +711,13 @@ def __call__( latent_dim = self.conditional_transformer.config.in_channels llm_features = self._encode_text( - inputs["token_ids"], inputs["text_position_ids"], inputs["indicator"] + inputs["token_ids"], + inputs["text_position_ids"], + inputs["text_indicator"], + max_text_tokens, + ) + llm_features = _append_image_feature_zeros( + llm_features, num_image_tokens, self.diffusion_device ) # Negative branch is image-only (asymmetric CFG) with zeroed conditioning. @@ -561,10 +729,10 @@ def __call__( num_image_tokens, llm_features.shape[-1], dtype=llm_features.dtype, - device=self.device, + device=self.diffusion_device, ) - generator = torch.Generator(device=self.device) + generator = torch.Generator(device=self.diffusion_device) if seed is not None: generator.manual_seed(seed) z = torch.randn( # type: ignore[call-overload] @@ -572,7 +740,7 @@ def __call__( num_image_tokens, latent_dim, dtype=torch.float32, - device=self.device, + device=self.diffusion_device, generator=generator, ) @@ -581,13 +749,15 @@ def __call__( max_text_tokens, latent_dim, dtype=torch.float32, - device=self.device, + device=self.diffusion_device, ) for i in range(num_steps - 1, -1, -1): t_val = float(schedule(step_intervals[i + 1].unsqueeze(0)).item()) s_val = float(schedule(step_intervals[i].unsqueeze(0)).item()) - t = torch.full((batch_size,), t_val, dtype=torch.float32, device=self.device) + t = torch.full( + (batch_size,), t_val, dtype=torch.float32, device=self.diffusion_device + ) pos_z = torch.cat([text_z_padding, z], dim=1) pos_out = self.conditional_transformer( @@ -614,6 +784,31 @@ def __call__( delta = s_val - t_val z = z + v * delta + # Do not retain the final step's branch outputs while the VAE allocates + # its decode workspace. These tensors are recreated on every step. + del pos_z, pos_out, pos_v, neg_v, v, t + + # Full text-plus-image conditioning is needed only by the diffusion models. + # At 1024x1024 the positive and negative float32 feature tensors occupy + # roughly 1.7 GiB together, enough to crowd out the VAE's peak workspace on + # a 20 GiB diffusion device if their references survive into _decode(). + del ( + llm_features, + neg_llm_features, + text_z_padding, + neg_position_ids, + neg_segment_ids, + neg_indicator, + inputs, + gw_per_step, + step_intervals, + generator, + ) + _synchronize_device(self.diffusion_device) + if self.diffusion_device.type == "cuda": + torch.cuda.empty_cache() + _log_device_memory([self.diffusion_device]) + return self._decode(z, grid_h=grid_h, grid_w=grid_w) # type: ignore[arg-type] def _decode(self, z: torch.Tensor, *, grid_h: int, grid_w: int) -> list[Image.Image]: diff --git a/tests/test_split_devices.py b/tests/test_split_devices.py new file mode 100644 index 0000000..3f789fa --- /dev/null +++ b/tests/test_split_devices.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +import unittest +import warnings +from unittest import mock + +import torch + +import ideogram4.pipeline_ideogram4 as pipeline_module +from ideogram4.constants import LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR +from ideogram4.pipeline_ideogram4 import ( + Ideogram4Pipeline, + Ideogram4PipelineConfig, + _append_image_feature_zeros, + _move_tensor_to_device, +) + + +def _bare_pipeline( + *, + diffusion_device: torch.device | None = None, + text_device: torch.device | None = None, +) -> Ideogram4Pipeline: + diffusion_device = diffusion_device or torch.device("cpu") + text_device = text_device or torch.device("cpu") + pipeline = Ideogram4Pipeline.__new__(Ideogram4Pipeline) + pipeline.config = Ideogram4PipelineConfig() + pipeline.diffusion_device = diffusion_device + pipeline.text_device = text_device + pipeline.device = diffusion_device + return pipeline + + +class SplitDeviceTests(unittest.TestCase): + def test_constructor_defaults_text_device_to_diffusion_device(self) -> None: + device = torch.device("cuda:0") + shift = mock.Mock() + scale = mock.Mock() + with ( + mock.patch.object( + pipeline_module, + "get_latent_norm", + return_value=(shift, scale), + ), + mock.patch.object(pipeline_module, "CaptionVerifier"), + ): + pipeline = Ideogram4Pipeline( + conditional_transformer=mock.Mock(), + unconditional_transformer=mock.Mock(), + text_encoder=mock.Mock(), + text_tokenizer=mock.Mock(), + autoencoder=mock.Mock(), + config=Ideogram4PipelineConfig(), + device=device, + dtype=torch.bfloat16, + ) + + self.assertEqual(pipeline.device, device) + self.assertEqual(pipeline.diffusion_device, device) + self.assertEqual(pipeline.text_device, device) + shift.to.assert_called_once_with(device) + scale.to.assert_called_once_with(device) + + def test_from_pretrained_routes_components_to_role_devices(self) -> None: + diffusion_device = torch.device("cpu") + text_device = torch.device("meta") + conditional = mock.Mock() + unconditional = mock.Mock() + + with ( + mock.patch.object( + pipeline_module, + "_load_indexed_or_single_state_dict", + side_effect=[{}, {}], + ), + mock.patch.object(pipeline_module, "hf_hub_download", return_value="vae"), + mock.patch.object( + pipeline_module, + "_build_transformer", + side_effect=[conditional, unconditional], + ) as build_transformer, + mock.patch.object( + pipeline_module, + "_load_qwen3_vl", + return_value=(mock.Mock(), mock.Mock()), + ) as load_qwen, + mock.patch.object( + pipeline_module, "_load_autoencoder", return_value=mock.Mock() + ) as load_autoencoder, + mock.patch.object( + pipeline_module, + "get_latent_norm", + return_value=(torch.zeros(1), torch.ones(1)), + ), + mock.patch.object(pipeline_module, "CaptionVerifier"), + mock.patch.object(pipeline_module, "_log_device_memory"), + ): + pipeline = Ideogram4Pipeline.from_pretrained( + device=diffusion_device, + text_device=text_device, + ) + + self.assertEqual(pipeline.diffusion_device, diffusion_device) + self.assertEqual(pipeline.text_device, text_device) + self.assertEqual(build_transformer.call_args_list[0].args[2], diffusion_device) + self.assertEqual(build_transformer.call_args_list[1].args[2], diffusion_device) + self.assertEqual(load_qwen.call_args.args[1], text_device) + self.assertEqual(load_autoencoder.call_args.args[1], diffusion_device) + + def test_build_inputs_assigns_tensors_to_role_devices(self) -> None: + pipeline = _bare_pipeline(text_device=torch.device("meta")) + pipeline._tokenize = mock.Mock(return_value=(torch.tensor([7, 8]), 2)) + + inputs = pipeline._build_inputs(["prompt"], height=16, width=16) + + for key in ("token_ids", "text_position_ids", "text_indicator"): + self.assertEqual(inputs[key].device, torch.device("meta")) + for key in ("position_ids", "segment_ids", "indicator"): + self.assertEqual(inputs[key].device, torch.device("cpu")) + + def test_encode_text_trims_trailing_image_positions(self) -> None: + pipeline = _bare_pipeline() + recorded: dict[str, int] = {} + + def fake_embeddings(token_ids, attention_mask, pos_2d): + recorded["sequence_length"] = token_ids.shape[1] + hidden = torch.ones(token_ids.shape[0], token_ids.shape[1], 2) + return [hidden for _ in pipeline_module.QWEN3_VL_ACTIVATION_LAYERS] + + pipeline._get_qwen3_vl_embeddings = fake_embeddings + max_text_tokens = 37 + total_length = max_text_tokens + 4096 + token_ids = torch.zeros(1, total_length, dtype=torch.long) + position_ids = torch.zeros(1, total_length, 3, dtype=torch.long) + text_indicator = torch.zeros(1, total_length, dtype=torch.long) + text_indicator[:, :max_text_tokens] = LLM_TOKEN_INDICATOR + + features = pipeline._encode_text( + token_ids, position_ids, text_indicator, max_text_tokens + ) + + self.assertEqual(recorded["sequence_length"], max_text_tokens) + self.assertEqual(features.shape[:2], (1, max_text_tokens)) + + def test_append_image_feature_zeros_restores_full_shape(self) -> None: + text_features = torch.randn(2, 5, 7) + features = _append_image_feature_zeros( + text_features, num_image_tokens=11, destination=torch.device("cpu") + ) + + self.assertEqual(features.shape, (2, 16, 7)) + torch.testing.assert_close(features[:, :5], text_features) + self.assertEqual(torch.count_nonzero(features[:, 5:]).item(), 0) + + def test_left_padding_stays_in_common_text_block(self) -> None: + pipeline = _bare_pipeline() + tokenized = { + "short": (torch.tensor([11, 12]), 2), + "long": (torch.tensor([21, 22, 23, 24]), 4), + } + pipeline._tokenize = mock.Mock(side_effect=lambda prompt: tokenized[prompt]) + + inputs = pipeline._build_inputs(["short", "long"], height=16, width=32) + max_text_tokens = inputs["max_text_tokens"] + self.assertEqual(max_text_tokens, 4) + torch.testing.assert_close(inputs["token_ids"][0, :4], torch.tensor([0, 0, 11, 12])) + torch.testing.assert_close( + inputs["token_ids"][1, :4], torch.tensor([21, 22, 23, 24]) + ) + torch.testing.assert_close( + inputs["text_indicator"][0, :4], + torch.tensor([0, 0, LLM_TOKEN_INDICATOR, LLM_TOKEN_INDICATOR]), + ) + self.assertTrue( + torch.all(inputs["indicator"][:, max_text_tokens:] == OUTPUT_IMAGE_INDICATOR) + ) + + def test_copy_fallback_stages_through_cpu(self) -> None: + destination = torch.device("cuda:1") + + class FakeTensor: + def __init__(self, device: str, *, fail_direct: bool = False): + self.device = torch.device(device) + self.fail_direct = fail_direct + + def to(self, device): + resolved = torch.device(device) + if self.fail_direct and resolved == destination: + raise RuntimeError("peer copy unavailable") + return FakeTensor(str(resolved)) + + source = FakeTensor("cuda:0", fail_direct=True) + with warnings.catch_warnings(record=True) as caught: + moved = _move_tensor_to_device(source, destination) + + self.assertEqual(moved.device, destination) + self.assertEqual(len(caught), 1) + self.assertIn("staging through CPU", str(caught[0].message)) + + def test_copy_fallback_does_not_swallow_staging_errors(self) -> None: + destination = torch.device("cuda:1") + + class BrokenTensor: + device = torch.device("cuda:0") + + def to(self, device): + if torch.device(device) == destination: + raise RuntimeError("peer copy unavailable") + raise ValueError("CPU staging failed") + + with ( + self.assertRaisesRegex(ValueError, "CPU staging failed"), + warnings.catch_warnings(), + ): + warnings.simplefilter("ignore") + _move_tensor_to_device(BrokenTensor(), destination) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_validate_comfyui_api.py b/tests/test_validate_comfyui_api.py new file mode 100644 index 0000000..dbd1965 --- /dev/null +++ b/tests/test_validate_comfyui_api.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import unittest + +from scripts.validate_comfyui_api import _saved_images, _workflow + + +class ComfyUIValidationTests(unittest.TestCase): + def test_workflow_uses_local_caption_and_acceptance_settings(self) -> None: + caption = '{"high_level_description":"test"}' + + workflow = _workflow(caption) + + self.assertEqual(workflow["1"]["class_type"], "Ideogram4PipelineLoader") + self.assertEqual(workflow["1"]["inputs"]["model_weights"], "4.0 NF4") + self.assertEqual(workflow["2"]["class_type"], "Ideogram4Generate") + self.assertEqual(workflow["2"]["inputs"]["pipeline"], ["1", 0]) + self.assertEqual(workflow["2"]["inputs"]["prompt"], caption) + self.assertEqual(workflow["2"]["inputs"]["width"], 1024) + self.assertEqual(workflow["2"]["inputs"]["height"], 1024) + self.assertEqual(workflow["2"]["inputs"]["sampler_preset"], "4.0 Default 20") + self.assertEqual(workflow["3"]["class_type"], "SaveImage") + self.assertEqual(workflow["3"]["inputs"]["images"], ["2", 0]) + + def test_saved_images_reads_save_node_output(self) -> None: + expected = [{"filename": "result.png", "subfolder": "", "type": "output"}] + + images = _saved_images({"outputs": {"3": {"images": expected}}}) + + self.assertEqual(images, expected) + + +if __name__ == "__main__": + unittest.main()