diff --git a/llada_gguf_adapter.py b/llada_gguf_adapter.py index ff82953..ab8e1c1 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. 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" @@ -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 @@ -896,7 +917,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 +925,31 @@ 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) + # 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"]: @@ -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) @@ -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") diff --git a/nodes.py b/nodes.py index 95803c9..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", "cpu"], + ["sequential_cpu_offload", "model_cpu_offload", "cuda", "gpu", "gpu_full", "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 ()