From a79e42a8cc9d7ffc8cd5e238531eebe55fa6bec4 Mon Sep 17 00:00:00 2001 From: gaius-codius <206332531+gaius-codius@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:13:16 +1000 Subject: [PATCH 1/2] Add a device-agnostic gpu offload mode for Intel XPU. The cuda branch was hard-coded to cuda:0 and refused to run on this Arc stack, so LLaDA stayed on CPU. gpu uses ComfyUI's selected device and keeps the same INT8 staging / aux-swap lifecycle. --- llada_gguf_adapter.py | 89 +++++++++++++++++++++++++------------------ nodes.py | 38 +++++++++--------- 2 files changed, 69 insertions(+), 58 deletions(-) diff --git a/llada_gguf_adapter.py b/llada_gguf_adapter.py index ff82953..d755409 100644 --- a/llada_gguf_adapter.py +++ b/llada_gguf_adapter.py @@ -19,6 +19,24 @@ log = logging.getLogger("ComfyUI-LLaDA-Image") +def _gpu_device() -> torch.device: + """The accelerator ComfyUI selected. No guessing: outside Comfy, or with + ONEAPI_DEVICE_SELECTOR unset, torch.device("xpu:0") is the B580, not the B60.""" + import comfy.model_management as mm # let an ImportError surface + dev = mm.get_torch_device() # on Intel: torch.device("xpu", torch.xpu.current_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" @@ -854,39 +872,38 @@ 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"): + # 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. # # 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.") - 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 @@ -896,7 +913,7 @@ 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) @@ -904,23 +921,22 @@ def _move(module, 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) + _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"]: @@ -938,35 +954,33 @@ def _release_aux_before_denoise(_module, _args): 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) @@ -975,7 +989,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") diff --git a/nodes.py b/nodes.py index 95803c9..499905f 100644 --- a/nodes.py +++ b/nodes.py @@ -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", "cpu"], {"default": "sequential_cpu_offload"}, ), "vae_tiling": (["On", "Auto", "Off"], {"default": "On"}), @@ -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", @@ -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, @@ -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 () From cd9e15b0f139dca5011b97bc3b4e4c651882dc95 Mon Sep 17 00:00:00 2001 From: gaius-codius <206332531+gaius-codius@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:13:16 +1000 Subject: [PATCH 2/2] Add gpu_full offload so the GGUF text encoder can run on the accelerator. gpu keeps prompt encode on the host. gpu_full moves registered encoder params so dequant temps land in VRAM; mmap'd qweights stay on the CPU. --- llada_gguf_adapter.py | 28 ++++++++++++++++++++++------ nodes.py | 2 +- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/llada_gguf_adapter.py b/llada_gguf_adapter.py index d755409..ab8e1c1 100644 --- a/llada_gguf_adapter.py +++ b/llada_gguf_adapter.py @@ -20,10 +20,10 @@ def _gpu_device() -> torch.device: - """The accelerator ComfyUI selected. No guessing: outside Comfy, or with - ONEAPI_DEVICE_SELECTOR unset, torch.device("xpu:0") is the B580, not the B60.""" + """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() # on Intel: torch.device("xpu", torch.xpu.current_device()) + 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 @@ -872,12 +872,15 @@ 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 in ("cuda", "gpu"): + 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 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 @@ -885,6 +888,7 @@ def build_llada_pipeline(selection: dict, pipeline_cls, config_dir: Path): 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" log.info("LLaDA GPU (%s): moving transformer non-quantized parameters to %s", gpu, gpu) transformer.to(device=gpu, dtype=dtype) @@ -927,6 +931,15 @@ def _begin_generation(generation_mode): _move(transformer, gpu) # Prompt preprocessing is required for BOTH Base/Turbo and T2I/Edit. + # 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) @@ -945,12 +958,15 @@ 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: diff --git a/nodes.py b/nodes.py index 499905f..2bb51d6 100644 --- a/nodes.py +++ b/nodes.py @@ -239,7 +239,7 @@ def INPUT_TYPES(cls): "vae": (vaes,), "dtype": (["bfloat16", "float16", "float32"], {"default": "bfloat16"}), "offload": ( - ["sequential_cpu_offload", "model_cpu_offload", "cuda", "gpu", "cpu"], + ["sequential_cpu_offload", "model_cpu_offload", "cuda", "gpu", "gpu_full", "cpu"], {"default": "sequential_cpu_offload"}, ), "vae_tiling": (["On", "Auto", "Off"], {"default": "On"}),