Skip to content

Add device-agnostic gpu / gpu_full offload for CUDA and XPU #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gaius-codius wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add device-agnostic gpu / gpu_full offload for CUDA and XPU #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a79e42a
Add a device-agnostic gpu offload mode for Intel XPU.
gaius-codius Sep 17, 2026
cd9e15b
Add gpu_full offload so the GGUF text encoder can run on the accelera…
gaius-codius Sep 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 70 additions & 39 deletions llada_gguf_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,24 @@
log = logging.getLogger("ComfyUI-LLaDA-Image")


def _gpu_device() -> torch.device:
"""The accelerator ComfyUI selected. Do not guess cuda:0/xpu:0: on a
multi-GPU host the default index may not be the card Comfy is using."""
import comfy.model_management as mm # let an ImportError surface
dev = mm.get_torch_device()
if dev.type not in ("cuda", "xpu"):
raise RuntimeError(f"LLaDA gpu mode needs a cuda/xpu device; ComfyUI selected {dev}.")
return dev


def _gpu_empty_cache(dev: torch.device) -> None:
if dev.type == "cuda" and torch.cuda.is_available():
torch.cuda.empty_cache()
elif dev.type == "xpu" and hasattr(torch, "xpu"):
torch.xpu.synchronize(dev) # without this, XPU often does not reclaim the
torch.xpu.empty_cache() # denoise activations before the VAE swap


def _load_city96_modules():
"""Load City96 ComfyUI-GGUF internals without executing its node __init__ again."""
alias = "_llada_city96_gguf"
Expand Down Expand Up @@ -854,39 +872,42 @@ def build_llada_pipeline(selection: dict, pipeline_cls, config_dir: Path):
pipe.enable_sequential_cpu_offload()
elif offload == "model_cpu_offload":
pipe.enable_model_cpu_offload()
elif offload == "cuda":
# Custom INT8 CUDA mode:
elif offload in ("cuda", "gpu", "gpu_full"):
# Custom INT8 GPU mode:
# - Keep LazyINT8Linear qweights/scales CPU-resident and stage them on demand.
# - Move the normal transformer parameters AND the auxiliary neural modules
# used before denoising to CUDA.
# used before denoising to the accelerator ComfyUI selected.
#
# gpu_full also places the GGUF text encoder's registered params on the
# GPU so prompt encode dequants there (mmap'd qweights stay on the host).
#
# This matters especially for native editing. Before the progress bar is
# created, the official pipeline runs QueryFormer/text projection, SigVQ
# reference-image encoding, and VAE reference-latent encoding. Leaving
# SigVQ/VAE on CPU can make an edit appear hung for hours.
if not torch.cuda.is_available():
gpu = _gpu_device()
if offload == "cuda" and gpu.type != "cuda":
raise RuntimeError("LLaDA CUDA mode selected but CUDA is not available.")
move_text_encoder = offload == "gpu_full"

cuda_device = torch.device("cuda:0")

log.info("LLaDA CUDA: moving transformer non-quantized parameters to cuda:0")
transformer.to(device=cuda_device, dtype=dtype)
log.info("LLaDA CUDA: transformer core ready; INT8 matrices remain lazy on CPU")
log.info("LLaDA GPU (%s): moving transformer non-quantized parameters to %s", gpu, gpu)
transformer.to(device=gpu, dtype=dtype)
log.info("LLaDA GPU (%s): transformer core ready; INT8 matrices remain lazy on CPU", gpu)

# Match the official all-CUDA pipeline for the auxiliary compute modules,
# but do it explicitly so our unregistered INT8 matrices are NOT eagerly
# migrated by DiffusionPipeline.to().
log.info("LLaDA CUDA: moving QueryFormer to cuda:0")
queryformer.to(device=cuda_device, dtype=dtype)
log.info("LLaDA GPU (%s): moving QueryFormer to %s", gpu, gpu)
queryformer.to(device=gpu, dtype=dtype)

log.info("LLaDA CUDA: moving text projection to cuda:0")
text_projection.to(device=cuda_device, dtype=dtype)
log.info("LLaDA GPU (%s): moving text projection to %s", gpu, gpu)
text_projection.to(device=gpu, dtype=dtype)

log.info("LLaDA CUDA: moving SigVQ to cuda:0 (required for fast native editing)")
sigvq.to(device=cuda_device, dtype=dtype)
log.info("LLaDA GPU (%s): moving SigVQ to %s (required for fast native editing)", gpu, gpu)
sigvq.to(device=gpu, dtype=dtype)

log.info("LLaDA CUDA: moving VAE to cuda:0 (required for fast edit image encoding/decoding)")
vae.to(device=cuda_device, dtype=dtype)
log.info("LLaDA GPU (%s): moving VAE to %s (required for fast edit image encoding/decoding)", gpu, gpu)
vae.to(device=gpu, dtype=dtype)

# Explicit reusable phase manager. nodes.py calls _rebel_begin_generation()
# before EVERY T2I/edit invocation. The transformer pre-hook then releases
Expand All @@ -896,31 +917,39 @@ def build_llada_pipeline(selection: dict, pipeline_cls, config_dir: Path):
def _move(module, device):
if module is None:
return
if str(device).startswith("cuda"):
if torch.device(device).type != "cpu":
module.to(device=torch.device(device), dtype=getattr(module, "dtype", dtype))
else:
module.to(device=device)

def _begin_generation(generation_mode):
phase["active"] = True
phase["released"] = False
cuda_device = "cuda:0"

# Restore the transformer core for every run. Lazy INT8 matrices are
# plain attrs, so they remain CPU-resident/staged on demand.
_move(transformer, cuda_device)
_move(transformer, gpu)

# Prompt preprocessing is required for BOTH Base/Turbo and T2I/Edit.
_move(queryformer, cuda_device)
_move(text_projection, cuda_device)
# gpu_full: GGUF qweights are unregistered attrs, so .to() only moves
# norms/gates; lazy layers then dequantise onto input_ids.device.
if move_text_encoder:
_move(text_encoder, gpu)
log.info(
"LLaDA GPU (%s): text encoder device is %s (GGUF matrices stay mmap'd on CPU)",
gpu,
text_encoder.device,
)
_move(queryformer, gpu)
_move(text_projection, gpu)

# SigVQ is only needed for VQ/editing. VAE encode is only needed for edit.
if generation_mode in ("vq", "editing"):
_move(sigvq, cuda_device)
_move(sigvq, gpu)
if generation_mode == "editing":
_move(vae, cuda_device)
_move(vae, gpu)

log.info("LLaDA CUDA phase: preprocessing ready for %s", generation_mode)
log.info("LLaDA GPU (%s) phase: preprocessing ready for %s", gpu, generation_mode)

def _release_aux_before_denoise(_module, _args):
if not phase["active"] or phase["released"]:
Expand All @@ -929,44 +958,45 @@ def _release_aux_before_denoise(_module, _args):

# Conditioning has already been built by the time transformer.forward
# is entered. Release aux VRAM so INT8 denoising has the full budget.
for _name, _component in (
release = [
("queryformer", queryformer),
("text_projection", text_projection),
("sigvq", sigvq),
("vae", vae),
):
]
if move_text_encoder:
release.insert(0, ("text_encoder", text_encoder))
for _name, _component in release:
try:
_component.to("cpu")
except Exception as _exc:
log.warning("LLaDA CUDA phase: could not release %s: %s", _name, _exc)
log.warning("LLaDA GPU (%s) phase: could not release %s: %s", gpu, _name, _exc)

if torch.cuda.is_available():
torch.cuda.empty_cache()
log.info("LLaDA CUDA phase: aux released; denoising begins")
_gpu_empty_cache(gpu)
log.info("LLaDA GPU (%s) phase: aux released; denoising begins", gpu)

transformer.register_forward_pre_hook(_release_aux_before_denoise)

# Decode wrapper fixes the post-step lifecycle correctly:
# official pipeline may have moved latents to CPU because VAE was released.
# Move transformer core OUT, move VAE + decode input BACK to CUDA, then decode.
# Move transformer core OUT, move VAE + decode input BACK to the GPU, then decode.
_original_vae_decode = vae.decode

def _cuda_decode(sample, *args, **kwargs):
if phase["active"]:
log.info("LLaDA CUDA phase: denoising complete; preparing CUDA VAE decode")
log.info("LLaDA GPU (%s) phase: denoising complete; preparing VAE decode", gpu)
try:
transformer.to("cpu")
finally:
if torch.cuda.is_available():
torch.cuda.empty_cache()
_gpu_empty_cache(gpu)

_move(vae, "cuda:0")
_move(vae, gpu)
if torch.is_tensor(sample):
sample = sample.to(device="cuda:0", dtype=vae.dtype)
sample = sample.to(device=gpu, dtype=vae.dtype)

result = _original_vae_decode(sample, *args, **kwargs)
phase["active"] = False
log.info("LLaDA CUDA phase: VAE decode complete")
log.info("LLaDA GPU (%s) phase: VAE decode complete", gpu)
return result

return _original_vae_decode(sample, *args, **kwargs)
Expand All @@ -975,7 +1005,8 @@ def _cuda_decode(sample, *args, **kwargs):
pipe._rebel_begin_generation = _begin_generation

log.info(
"LLaDA CUDA: reusable preprocess -> denoise -> decode phase manager ready"
"LLaDA GPU (%s): reusable preprocess -> denoise -> decode phase manager ready",
gpu,
)
elif offload == "cpu":
pipe.to("cpu")
Expand Down
38 changes: 17 additions & 21 deletions nodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,7 @@ def INPUT_TYPES(cls):
"vae": (vaes,),
"dtype": (["bfloat16", "float16", "float32"], {"default": "bfloat16"}),
"offload": (
["sequential_cpu_offload", "model_cpu_offload", "cuda", "cpu"],
["sequential_cpu_offload", "model_cpu_offload", "cuda", "gpu", "gpu_full", "cpu"],
{"default": "sequential_cpu_offload"},
),
"vae_tiling": (["On", "Auto", "Off"], {"default": "On"}),
Expand Down Expand Up @@ -296,20 +296,18 @@ def generate(
seed,
negative_prompt="",
):
# Diffusers creates the initial latent on pipeline._execution_device.
# Under CPU/model/sequential offload that device can be CPU even when
# CUDA exists, so a hard-coded CUDA Generator crashes randn_tensor.
execution_device = getattr(pipeline, "_execution_device", None)
if execution_device is None:
execution_device = getattr(pipeline, "device", torch.device("cpu"))
execution_device = torch.device(execution_device)
generator_device = "cuda" if execution_device.type == "cuda" else "cpu"
generator = torch.Generator(device=generator_device).manual_seed(int(seed))

# Latents are created on pipeline.transformer.device, not _execution_device.
# Under CPU/model/sequential offload that can be CPU even when a GPU
# exists, so a hard-coded CUDA Generator crashes randn_tensor.
begin_generation = getattr(pipeline, "_rebel_begin_generation", None)
if callable(begin_generation):
begin_generation("text")

gen_device = pipeline.transformer.device
if gen_device.type not in ("cuda", "xpu"):
gen_device = torch.device("cpu")
generator = torch.Generator(device=gen_device).manual_seed(int(seed))

result = pipeline(
prompt=prompt,
generation_mode="text",
Expand Down Expand Up @@ -394,22 +392,18 @@ def edit(
if image is None:
raise ValueError("LLaDA Image Edit requires a source IMAGE.")

# Diffusers creates the initial latent on pipeline._execution_device.
# Under CPU/model/sequential offload that device can be CPU even when
# CUDA exists, so a hard-coded CUDA Generator crashes randn_tensor.
execution_device = getattr(pipeline, "_execution_device", None)
if execution_device is None:
execution_device = getattr(pipeline, "device", torch.device("cpu"))
execution_device = torch.device(execution_device)
generator_device = "cuda" if execution_device.type == "cuda" else "cpu"
generator = torch.Generator(device=generator_device).manual_seed(int(seed))

source_image = _comfy_image_to_pil(image)

# Latents are created on pipeline.transformer.device, not _execution_device.
begin_generation = getattr(pipeline, "_rebel_begin_generation", None)
if callable(begin_generation):
begin_generation("editing")

gen_device = pipeline.transformer.device
if gen_device.type not in ("cuda", "xpu"):
gen_device = torch.device("cpu")
generator = torch.Generator(device=gen_device).manual_seed(int(seed))

result = pipeline(
prompt=prompt,
image=source_image,
Expand Down Expand Up @@ -454,6 +448,8 @@ def unload(self, pipeline):
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
return ()


Expand Down