llada t2i

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -16,6 +16,10 @@ All-about-optimizations:
 - **Models**
   - [Anima 2.9B Preview v1](https://huggingface.co/yeoj34760/Anima-2.9B)  
     expanded version of Anima 2B  
+  - [inclusionAI LLaDA-Image](https://huggingface.co/inclusionAI/LLaDA-Image) in *base* and *turbo* variants  
+    LLaDA-Image is a 6.5B transformer with massive 16.3B fully-custom MoE text-encoder and optional 1.3B SigVQ conditioning model  
+    with support for text-to-image, vq-conditioned text-to-image and image-editing workflows  
+    *note* model is extremely quantization sensitive so minimum allowed quant type is `uint8`  
 - **LoRA**
   - *TODO*: see [LoRA docs](https://vladmandic.github.io/sdnext-docs/LoRA) for all of the improvements and usage instructions  
     *note*: lora now has its own settings section in *settings -> lora*

modules/modeldata.py

@@ -119,6 +119,8 @@ def get_model_type(pipe):
         model_type = 'longcat'
     elif 'GlmImage' in name:
         model_type = 'glmimage'
+    elif 'LLaDAImage' in name:
+        model_type = 'lladaimage'
     elif 'Step1XEdit' in name:
         model_type = 'step1x_edit'
     elif 'JoyImageEdit' in name:

modules/sd_detect.py

@@ -167,6 +167,8 @@ def guess_by_name(fn, current_guess):
         new_guess = 'OvisImage'
     elif 'glm-image' in fn.lower():
         new_guess = 'GLMImage'
+    elif 'llada' in fn.lower():
+        new_guess = 'LLaDAImage'
     elif 'sdxs-1b' in fn.lower():
         new_guess = 'SDXS'
     elif 'step1x-edit' in fn.lower():

modules/sd_models.py

@@ -55,6 +55,7 @@ pipe_switch_task_exclude = [
     'Kandinsky5I2IPipeline',
     'GoogleNanoBananaPipeline',
     'Step1XEditPipeline',
+    'LLaDAImagePipeline',
     'BooguImagePipeline',
     'BooguImageTurboPipeline',
 ]
@@ -610,6 +611,10 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf
             from pipelines.model_glm import load_glm_image
             sd_model = load_glm_image(checkpoint_info, diffusers_load_config)
             allow_post_quant = False
+        elif model_type in ['LLaDAImage']:
+            from pipelines.model_llada import load_llada_image
+            sd_model = load_llada_image(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['SDXS']:
             from pipelines.model_sdxs import load_sdxs
             sd_model = load_sdxs(checkpoint_info, diffusers_load_config)

modules/vae/sd_vae_taesd.py

@@ -73,7 +73,7 @@ def get_model(model_cls, variant=None):
     elif model_cls in {'f1', 'h1', 'zimage', 'lumina2', 'chroma', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4', 'ultraflux'}:
         model_cls = 'f1'
         variant = 'TAE FLUX.1'
-    elif model_cls in {'f2', 'ernieimage', 'lens', 'ideogram4'}:
+    elif model_cls in {'f2', 'ernieimage', 'lens', 'ideogram4', 'lladaimage'}:
         model_cls = 'f2'
         variant = 'TAE FLUX.2'
     elif model_cls in {'sd3'}:

pipelines/llada/__init__.py

@@ -0,0 +1,8 @@
+from .pipeline_llada_image import LLaDAImagePipeline
+from .pipeline_output import LLaDAImagePipelineOutput
+from .transformer_llada_image import LLaDAImageQueryFormerModel
+from .transformer_llada_image import LLaDAImageSigVQModel
+from .transformer_llada_image import LLaDAImageTextProjectionModel
+from .transformer_llada_image import LLaDAImageTransformer2DModel
+
+__all__ = ['LLaDAImagePipeline', 'LLaDAImagePipelineOutput', 'LLaDAImageQueryFormerModel', 'LLaDAImageSigVQModel', 'LLaDAImageTextProjectionModel', 'LLaDAImageTransformer2DModel']

pipelines/llada/configuration_llada2uni_moe.py

@@ -0,0 +1,133 @@
+# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""LLaDA2 MoE model configuration."""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class LLaDA2MoeConfig(PretrainedConfig):
+    r"""
+    Configuration class for the LLaDA2 MoE model.
+
+    ```python
+    >>> from configuration_llada2uni_moe import LLaDA2MoeConfig
+    >>> config = LLaDA2MoeConfig()
+    ```
+    """
+
+    # Keep the original value because it selects the fused-expert implementation.
+    model_type = "llada2_moe_veomni"
+
+    def __init__(
+        self,
+        vocab_size=30592,
+        hidden_size=1024,
+        intermediate_size=None,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_key_value_heads=0,
+        head_dim=None,
+        hidden_act="silu",
+        use_qkv_bias=False,
+        use_qk_norm=True,
+        use_bias=True,
+        rms_norm_eps=1e-05,
+        tie_word_embeddings=False,
+        attention_dropout=0.1,
+        initializer_range=0.02,
+        max_position_embeddings=16384,
+        rope_theta=10000.0,
+        rope_parameters=None,
+        rope_scaling=None,
+        partial_rotary_factor=0.5,
+        use_cache=True,
+        sliding_window=None,
+        pad_token_id=126081,
+        # Image
+        image_token_offset=157184,
+        # MoE
+        num_experts=16,
+        num_shared_experts=0,
+        num_experts_per_tok=2,
+        n_group=8,
+        topk_group=4,
+        routed_scaling_factor=2.5,
+        moe_router_enable_expert_bias=True,
+        norm_topk_prob=True,
+        router_dtype="fp32",
+        score_function="sigmoid",
+        moe_intermediate_size=None,
+        first_k_dense_replace=0,
+        output_router_logits=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim or hidden_size // num_attention_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_qk_norm = use_qk_norm
+        self.use_bias = use_bias
+        self.rms_norm_eps = rms_norm_eps
+        self.attention_dropout = attention_dropout
+        self.initializer_range = initializer_range
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.partial_rotary_factor = partial_rotary_factor
+        self.use_cache = use_cache
+        self.sliding_window = sliding_window
+
+        # Image token offset: VQ codebook indices are shifted by this amount in the vocabulary
+        self.image_token_offset = image_token_offset
+
+        # RoPE parameters dict — used by LLaDA2MoeRotaryEmbedding
+        if rope_parameters is None:
+            rope_parameters = {
+                "rope_type": "default",
+                "rope_theta": rope_theta,
+                "partial_rotary_factor": partial_rotary_factor,
+            }
+        self.rope_parameters = rope_parameters
+
+        # MoE
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.routed_scaling_factor = routed_scaling_factor
+        self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
+        self.norm_topk_prob = norm_topk_prob
+        self.router_dtype = router_dtype
+        self.score_function = score_function
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+
+        # FP8 quantization flag — set to True to use FP8Linear for experts
+        self.use_fp8_experts = kwargs.pop("use_fp8_experts", False)
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["LLaDA2MoeConfig"]

pipelines/llada/fused_moe_ops.py

@@ -0,0 +1,375 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Standalone, inference-only VeOmni v0.1.0 fused-MoE compatibility shim.
+
+This module preserves the ``veomni.ops.fused_moe_forward`` call signature used
+by VeOmni v0.1.0 while removing VeOmni's training, Expert Parallelism (EP), NPU,
+and Seed-kernel dependencies. It is intended for single-device inference only.
+
+The CUDA fast path uses a small Triton grouped-linear kernel. If Triton is not
+available, the tensors are not on CUDA, or ``LLADA_MOE_BACKEND=eager`` is set,
+the implementation falls back to ordinary PyTorch operations.
+
+Replace the original model-code import with, for example,
+``from .fused_moe_v010 import fused_moe_forward``.
+
+Derived from ByteDance-Seed/VeOmni v0.1.0.post1:
+https://github.com/ByteDance-Seed/VeOmni/tree/v0.1.0.post1
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+import torch.nn.functional as F
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError:  # The eager fallback does not require Triton.
+    triton = None
+    tl = None
+
+
+_SUPPORTED_TRITON_DTYPES = (torch.float16, torch.bfloat16)
+
+
+if triton is not None:
+
+    @triton.jit
+    def _grouped_linear_kernel(
+        input_ptr,
+        weight_ptr,
+        output_ptr,
+        expert_cumsum_ptr,
+        N: tl.constexpr,
+        K: tl.constexpr,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        BLOCK_K: tl.constexpr,
+    ):
+        """Compute per-expert ``input @ weight.T`` for contiguous tensors."""
+        block_m = tl.program_id(axis=0)
+        block_n = tl.program_id(axis=1)
+        expert = tl.program_id(axis=2)
+
+        expert_start = tl.load(expert_cumsum_ptr + expert - 1, mask=expert > 0, other=0)
+        expert_end = tl.load(expert_cumsum_ptr + expert)
+        expert_tokens = expert_end - expert_start
+
+        if block_m * BLOCK_M >= expert_tokens:
+            return
+
+        row_offsets = block_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        col_offsets = block_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        k_offsets = tl.arange(0, BLOCK_K)
+
+        input_ptrs = (
+            input_ptr
+            + (expert_start + row_offsets[:, None]) * K
+            + k_offsets[None, :]
+        )
+        weight_ptrs = (
+            weight_ptr
+            + expert * N * K
+            + col_offsets[None, :] * K
+            + k_offsets[:, None]
+        )
+
+        accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k_block in range(0, tl.cdiv(K, BLOCK_K)):
+            remaining_k = K - k_block * BLOCK_K
+            inputs = tl.load(
+                input_ptrs,
+                mask=(row_offsets[:, None] < expert_tokens) & (k_offsets[None, :] < remaining_k),
+                other=0.0,
+            )
+            weights = tl.load(
+                weight_ptrs,
+                mask=(col_offsets[None, :] < N) & (k_offsets[:, None] < remaining_k),
+                other=0.0,
+            )
+            accumulator += tl.dot(inputs, weights)
+            input_ptrs += BLOCK_K
+            weight_ptrs += BLOCK_K
+
+        output_ptrs = (
+            output_ptr
+            + (expert_start + row_offsets[:, None]) * N
+            + col_offsets[None, :]
+        )
+        tl.store(
+            output_ptrs,
+            accumulator,
+            mask=(row_offsets[:, None] < expert_tokens) & (col_offsets[None, :] < N),
+        )
+
+
+def _validate_inputs(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> None:
+    if num_experts <= 0:
+        raise ValueError(f"num_experts must be positive, got {num_experts}")
+    if torch.is_grad_enabled():
+        raise RuntimeError(
+            "This standalone fused_moe_forward is inference-only. Call it under "
+            "torch.no_grad() or torch.inference_mode()."
+        )
+    if hidden_states.ndim != 2:
+        raise ValueError(f"hidden_states must have shape [tokens, hidden], got {tuple(hidden_states.shape)}")
+    if routing_weights.ndim != 2 or selected_experts.shape != routing_weights.shape:
+        raise ValueError(
+            "routing_weights and selected_experts must have the same [tokens, top_k] shape, got "
+            f"{tuple(routing_weights.shape)} and {tuple(selected_experts.shape)}"
+        )
+    if routing_weights.shape[1] == 0:
+        raise ValueError("top_k must be positive")
+    if routing_weights.shape[0] != hidden_states.shape[0]:
+        raise ValueError("routing_weights and hidden_states must contain the same number of tokens")
+    if selected_experts.dtype not in (torch.int32, torch.int64):
+        raise TypeError(f"selected_experts must be int32 or int64, got {selected_experts.dtype}")
+    if fc1_1_weight.ndim != 3 or fc1_2_weight.ndim != 3 or fc2_weight.ndim != 3:
+        raise ValueError("expert weights must be rank-3 tensors")
+    if fc1_1_weight.shape != fc1_2_weight.shape:
+        raise ValueError("fc1_1_weight and fc1_2_weight must have identical shapes")
+
+    experts, intermediate_size, hidden_size = fc1_1_weight.shape
+    expected_fc2_shape = (experts, hidden_size, intermediate_size)
+    if experts != num_experts:
+        raise ValueError(f"num_experts={num_experts}, but the weights contain {experts} experts")
+    if hidden_states.shape[1] != hidden_size:
+        raise ValueError(f"hidden size is {hidden_states.shape[1]}, but the weights expect {hidden_size}")
+    if tuple(fc2_weight.shape) != expected_fc2_shape:
+        raise ValueError(f"fc2_weight must have shape {expected_fc2_shape}, got {tuple(fc2_weight.shape)}")
+    if selected_experts.numel():
+        # These scalar checks synchronize CUDA once, before launching harder-to-debug kernels.
+        min_expert = int(selected_experts.min().item())
+        max_expert = int(selected_experts.max().item())
+        if min_expert < 0 or max_expert >= num_experts:
+            raise ValueError(f"selected expert IDs must be in [0, {num_experts}), got [{min_expert}, {max_expert}]")
+
+    devices = {
+        hidden_states.device,
+        routing_weights.device,
+        selected_experts.device,
+        fc1_1_weight.device,
+        fc1_2_weight.device,
+        fc2_weight.device,
+    }
+    if len(devices) != 1:
+        raise ValueError(f"all inputs and weights must be on one device, got {sorted(map(str, devices))}")
+
+
+def _route_tokens(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Sort routed token copies by expert and return the inverse permutation."""
+    top_k = selected_experts.shape[1]
+    flat_experts = selected_experts.reshape(-1).to(torch.int64)
+    order = torch.argsort(flat_experts, stable=True)
+    sorted_hidden_states = hidden_states[torch.div(order, top_k, rounding_mode="floor")].contiguous()
+    sorted_routing_weights = routing_weights.reshape(-1)[order].contiguous()
+    tokens_per_expert = torch.bincount(flat_experts, minlength=num_experts)
+    expert_cumsum = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32).contiguous()
+    return sorted_hidden_states, sorted_routing_weights, expert_cumsum, order
+
+
+def _unroute_tokens(
+    sorted_outputs: torch.Tensor,
+    order: torch.Tensor,
+    num_tokens: int,
+    top_k: int,
+) -> torch.Tensor:
+    restored = torch.empty_like(sorted_outputs)
+    restored[order] = sorted_outputs
+    # VeOmni's v0.1.0 gather kernel accumulates the top-k outputs in FP32.
+    return restored.view(num_tokens, top_k, -1).sum(dim=1, dtype=torch.float32).to(sorted_outputs.dtype)
+
+
+def _grouped_linear_triton(
+    inputs: torch.Tensor,
+    weights: torch.Tensor,
+    expert_cumsum: torch.Tensor,
+) -> torch.Tensor:
+    if triton is None:  # pragma: no cover - guarded by the caller
+        raise RuntimeError("Triton is not available")
+    if not inputs.is_contiguous() or not weights.is_contiguous():
+        raise ValueError("the Triton path requires contiguous inputs and expert weights")
+
+    num_experts, output_size, input_size = weights.shape
+    if inputs.shape[1] != input_size:
+        raise ValueError(f"input width is {inputs.shape[1]}, but the weights expect {input_size}")
+
+    output = torch.empty((inputs.shape[0], output_size), dtype=inputs.dtype, device=inputs.device)
+    block_m, block_n, block_k = 128, 128, 32
+    grid = (
+        triton.cdiv(inputs.shape[0], block_m),
+        triton.cdiv(output_size, block_n),
+        num_experts,
+    )
+    with torch.cuda.device(inputs.device):
+        _grouped_linear_kernel[grid](
+            inputs,
+            weights,
+            output,
+            expert_cumsum,
+            N=output_size,
+            K=input_size,
+            BLOCK_M=block_m,
+            BLOCK_N=block_n,
+            BLOCK_K=block_k,
+            num_warps=8,
+            num_stages=3,
+        )
+    return output
+
+
+def _triton_moe_forward(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> torch.Tensor:
+    sorted_hidden, sorted_routing, expert_cumsum, order = _route_tokens(
+        num_experts, routing_weights, selected_experts, hidden_states
+    )
+    gate = _grouped_linear_triton(sorted_hidden, fc1_1_weight, expert_cumsum)
+    up = _grouped_linear_triton(sorted_hidden, fc1_2_weight, expert_cumsum)
+    intermediate = F.silu(gate) * up
+    intermediate.mul_(sorted_routing.unsqueeze(-1))
+    sorted_outputs = _grouped_linear_triton(intermediate.contiguous(), fc2_weight, expert_cumsum)
+    return _unroute_tokens(sorted_outputs, order, hidden_states.shape[0], selected_experts.shape[1])
+
+
+def _eager_moe_forward(
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> torch.Tensor:
+    sorted_hidden, sorted_routing, expert_cumsum, order = _route_tokens(
+        num_experts, routing_weights, selected_experts, hidden_states
+    )
+    expert_ends = expert_cumsum.to(device="cpu", dtype=torch.int64).tolist()
+    outputs: list[torch.Tensor] = []
+    start = 0
+    for expert, end in enumerate(expert_ends):
+        if end > start:
+            expert_inputs = sorted_hidden[start:end]
+            gate = F.linear(expert_inputs, fc1_1_weight[expert])
+            up = F.linear(expert_inputs, fc1_2_weight[expert])
+            intermediate = F.silu(gate) * up
+            intermediate.mul_(sorted_routing[start:end].unsqueeze(-1))
+            outputs.append(F.linear(intermediate, fc2_weight[expert]))
+        start = end
+
+    sorted_outputs = torch.cat(outputs, dim=0) if outputs else hidden_states.new_empty((0, hidden_states.shape[1]))
+    return _unroute_tokens(sorted_outputs, order, hidden_states.shape[0], selected_experts.shape[1])
+
+
+def fused_moe_forward(
+    module: torch.nn.Module,
+    num_experts: int,
+    routing_weights: torch.Tensor,
+    selected_experts: torch.Tensor,
+    hidden_states: torch.Tensor,
+    fc1_1_weight: torch.Tensor,
+    fc1_2_weight: torch.Tensor,
+    fc2_weight: torch.Tensor,
+) -> torch.Tensor:
+    """Run the VeOmni v0.1.0 split-weight MoE operation for inference.
+
+    ``module`` is retained for call-site compatibility. Like VeOmni's original
+    non-EP implementation, this function does not use it.
+
+    Set ``LLADA_MOE_BACKEND`` to ``auto`` (default), ``triton``, or ``eager``.
+    The ``triton`` setting fails loudly if its requirements are not met;
+    ``auto`` falls back to the PyTorch implementation.
+    """
+    del module
+    _validate_inputs(
+        num_experts,
+        routing_weights,
+        selected_experts,
+        hidden_states,
+        fc1_1_weight,
+        fc1_2_weight,
+        fc2_weight,
+    )
+
+    backend = os.getenv("LLADA_MOE_BACKEND", "auto").lower()
+    if backend not in {"auto", "triton", "eager"}:
+        raise ValueError(f"LLADA_MOE_BACKEND must be auto, triton, or eager; got {backend!r}")
+
+    compute_dtype = fc1_1_weight.dtype
+    if fc1_2_weight.dtype != compute_dtype or fc2_weight.dtype != compute_dtype:
+        raise TypeError("all expert weights must have the same dtype")
+    hidden_states = hidden_states.to(dtype=compute_dtype)
+    routing_weights = routing_weights.to(dtype=compute_dtype)
+
+    if hidden_states.shape[0] == 0:
+        return hidden_states
+
+    can_use_triton = (
+        triton is not None
+        and hidden_states.is_cuda
+        and compute_dtype in _SUPPORTED_TRITON_DTYPES
+        and fc1_1_weight.is_contiguous()
+        and fc1_2_weight.is_contiguous()
+        and fc2_weight.is_contiguous()
+    )
+    if backend == "triton" and not can_use_triton:
+        raise RuntimeError(
+            "The Triton backend requires Triton, CUDA tensors, contiguous expert weights, "
+            "and float16 or bfloat16 weights."
+        )
+    if backend != "eager" and can_use_triton:
+        return _triton_moe_forward(
+            num_experts,
+            routing_weights,
+            selected_experts,
+            hidden_states,
+            fc1_1_weight,
+            fc1_2_weight,
+            fc2_weight,
+        )
+    return _eager_moe_forward(
+        num_experts,
+        routing_weights,
+        selected_experts,
+        hidden_states,
+        fc1_1_weight,
+        fc1_2_weight,
+        fc2_weight,
+    )
+
+
+__all__ = ["fused_moe_forward"]

pipelines/llada/pipeline_llada_image.py

@@ -0,0 +1,687 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Callable
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.models import AutoencoderKLFlux2
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+
+from .transformer_llada_image import (
+    LLaDAImageQueryFormerModel,
+    LLaDAImageSigVQModel,
+    LLaDAImageTextProjectionModel,
+    LLaDAImageTransformer2DModel,
+)
+from .pipeline_output import LLaDAImagePipelineOutput
+
+
+logger = logging.get_logger(__name__)
+
+
+class LLaDAImagePipeline(DiffusionPipeline):
+    r"""
+    Pipeline for LLaDA-Image text-to-image generation, VQ-conditioned generation, and single-image editing.
+
+    Args:
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            Flow-matching scheduler used for denoising.
+        vae ([`AutoencoderKLFlux2`]):
+            Flux2 VAE used to encode reference images and decode generated latents.
+        text_encoder (`transformers.PreTrainedModel`):
+            LLaDA2 conditional-generation model. It must expose `get_input_embeddings()` and its language backbone as
+            `model`.
+        tokenizer (`transformers.PreTrainedTokenizerBase`):
+            Tokenizer paired with the LLaDA2 text encoder.
+        queryformer ([`LLaDAImageQueryFormerModel`]):
+            QueryFormer that refines the learnable generation queries.
+        text_projection ([`LLaDAImageTextProjectionModel`]):
+            Connector and projector that map LLaDA2 hidden states to denoiser caption features.
+        sigvq ([`LLaDAImageSigVQModel`]):
+            GLM SigVQ component that embeds MLLM-generated VQ tokens and encodes editing reference images.
+        transformer ([`LLaDAImageTransformer2DModel`]):
+            Denoising transformer.
+    """
+
+    model_cpu_offload_seq = "text_encoder->queryformer->text_projection->sigvq->transformer->vae"
+    _callback_tensor_inputs = ["latents", "noise_pred"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLFlux2,
+        text_encoder: PreTrainedModel,
+        tokenizer: PreTrainedTokenizerBase,
+        queryformer: LLaDAImageQueryFormerModel,
+        text_projection: LLaDAImageTextProjectionModel,
+        sigvq: LLaDAImageSigVQModel,
+        transformer: LLaDAImageTransformer2DModel,
+    ):
+        super().__init__()
+        self.register_modules(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            queryformer=queryformer,
+            text_projection=text_projection,
+            sigvq=sigvq,
+            transformer=transformer,
+        )
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if self.vae is not None else 8
+        self.latent_scale_factor = self.vae_scale_factor * 2
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.latent_scale_factor)
+
+    @classmethod
+    def from_pretrained(  # pylint: disable=arguments-differ
+        cls,
+        pretrained_model_name_or_path: str | Path,
+        *,
+        torch_dtype: torch.dtype | None = None,
+        device: torch.device | str | None = None,
+        cache_dir: str | None = None,
+        scheduler=None,
+        vae=None,
+        text_encoder=None,
+        tokenizer=None,
+        queryformer=None,
+        text_projection=None,
+        sigvq=None,
+        transformer=None,
+        **kwargs,
+    ) -> "LLaDAImagePipeline":
+        """Load all LLaDA-Image components from a converted model directory or Hugging Face repository.
+
+        This model stores a LLaDA2 text encoder that requires `trust_remote_code=True`, so its components are loaded
+        explicitly instead of relying on the generic Diffusers pipeline resolver.
+        """
+        model_path = Path(pretrained_model_name_or_path)
+        if not model_path.is_dir():
+            from huggingface_hub import snapshot_download
+
+            ignore_patterns = ["assets/**"]
+            if text_encoder is not None:
+                ignore_patterns.append("text_encoder/**")
+            if transformer is not None:
+                ignore_patterns.append("transformer/**")
+            model_path = Path(
+                snapshot_download(
+                    repo_id=str(pretrained_model_name_or_path),
+                    cache_dir=cache_dir,
+                    ignore_patterns=ignore_patterns or None,
+                )
+            )
+
+        if not (model_path / "model_index.json").is_file():
+            raise ValueError(
+                "Expected a converted LLaDA-Image model directory containing `model_index.json`, got "
+                f"{model_path}."
+            )
+
+        if scheduler is None:
+            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_path / "scheduler", cache_dir=cache_dir)
+        if vae is None:
+            vae = AutoencoderKLFlux2.from_pretrained(model_path / "vae", torch_dtype=torch_dtype, cache_dir=cache_dir)
+        if text_encoder is None:
+            text_encoder_kwargs = {"dtype": torch_dtype, "trust_remote_code": True, "cache_dir": cache_dir}
+            if device is not None:
+                text_encoder_kwargs["device_map"] = {"": device}
+            text_encoder = AutoModel.from_pretrained(model_path / "text_encoder", **text_encoder_kwargs)
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(model_path / "tokenizer", cache_dir=cache_dir)
+        if queryformer is None:
+            queryformer = LLaDAImageQueryFormerModel.from_pretrained(
+                model_path / "queryformer", torch_dtype=torch_dtype, cache_dir=cache_dir
+            )
+        if text_projection is None:
+            text_projection = LLaDAImageTextProjectionModel.from_pretrained(
+                model_path / "text_projection", torch_dtype=torch_dtype, cache_dir=cache_dir
+            )
+        if sigvq is None:
+            sigvq = LLaDAImageSigVQModel.from_pretrained(
+                model_path / "sigvq", torch_dtype=torch_dtype, cache_dir=cache_dir
+            )
+        if transformer is None:
+            transformer = LLaDAImageTransformer2DModel.from_pretrained(
+                model_path / "transformer", torch_dtype=torch_dtype, cache_dir=cache_dir
+            )
+
+        if device is not None:
+            vae = vae.to(device)
+            queryformer = queryformer.to(device)
+            text_projection = text_projection.to(device)
+            sigvq = sigvq.to(device)
+            transformer = transformer.to(device)
+
+        return cls(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            queryformer=queryformer,
+            text_projection=text_projection,
+            sigvq=sigvq,
+            transformer=transformer,
+        )
+
+    @property
+    def guidance_scale(self) -> float:
+        return self._guidance_scale
+
+    @property
+    def num_timesteps(self) -> int:
+        return self._num_timesteps
+
+    @staticmethod
+    def _patchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        batch_size, channels, height, width = latents.shape
+        latents = latents.reshape(batch_size, channels, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 1, 3, 5, 2, 4)
+        return latents.reshape(batch_size, channels * 4, height // 2, width // 2)
+
+    @staticmethod
+    def _unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        batch_size, channels, height, width = latents.shape
+        latents = latents.reshape(batch_size, channels // 4, 2, 2, height, width)
+        latents = latents.permute(0, 1, 4, 2, 5, 3)
+        return latents.reshape(batch_size, channels // 4, height * 2, width * 2)
+
+    def _encode_text(
+        self,
+        prompts: list[str],
+        max_sequence_length: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        formatted_prompts = [
+            "<role>HUMAN</role> Generate an image.\n<role>ASSISTANT</role>\n<IMAGE1>"
+            if prompt is None
+            else f"<role>HUMAN</role> Generate an image: {prompt.strip()}\n<role>ASSISTANT</role>\n<IMAGE1>"
+            for prompt in prompts
+        ]
+        text_inputs = self.tokenizer(
+            formatted_prompts,
+            add_special_tokens=True,
+            padding=True,
+            truncation=True,
+            max_length=max_sequence_length,
+            return_tensors="pt",
+        )
+        input_ids = text_inputs.input_ids.to(self.text_encoder.device)
+        attention_mask = text_inputs.attention_mask.to(input_ids.device).bool()
+        inputs_embeds = self.text_encoder.get_input_embeddings()(input_ids)
+        text_encoder_device = inputs_embeds.device
+        attention_mask = attention_mask.to(text_encoder_device)
+
+        query_embeds = self.queryformer(
+            inputs_embeds.to(device=self.queryformer.device, dtype=self.queryformer.dtype),
+            attention_mask.to(self.queryformer.device),
+        ).query_embeds.to(device=text_encoder_device, dtype=inputs_embeds.dtype)
+        text_length = inputs_embeds.shape[1]
+        inputs_embeds = torch.cat([inputs_embeds, query_embeds], dim=1)
+        attention_mask = torch.cat(
+            [attention_mask, attention_mask.new_ones(attention_mask.shape[0], query_embeds.shape[1])],
+            dim=1,
+        )
+        position_ids = attention_mask.long().cumsum(dim=1) - 1
+        position_ids.masked_fill_(position_ids < 0, 0)
+
+        mask_value = torch.finfo(inputs_embeds.dtype).min
+        backbone_attention_mask = attention_mask[:, None, None, :].expand(-1, 1, attention_mask.shape[1], -1)
+        backbone_attention_mask = torch.where(
+            backbone_attention_mask,
+            torch.zeros((), dtype=inputs_embeds.dtype, device=text_encoder_device),
+            torch.full((), mask_value, dtype=inputs_embeds.dtype, device=text_encoder_device),
+        )
+        backbone_attention_mask[:, :, :text_length, text_length:] = mask_value
+
+        hidden_states = self.text_encoder.model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=backbone_attention_mask,
+            position_ids=position_ids,
+            return_dict=True,
+        ).last_hidden_state
+        prompt_embeds = self.text_projection(
+            hidden_states.to(device=self.text_projection.device, dtype=self.text_projection.dtype)
+        ).hidden_states
+        return prompt_embeds, attention_mask.to(prompt_embeds.device)
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str] | None,
+        negative_prompt: str | list[str] | None = None,
+        do_classifier_free_guidance: bool = True,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_attention_mask: torch.Tensor | None = None,
+        max_sequence_length: int = 2048,
+        device: torch.device | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        device = device or self._execution_device
+
+        if prompt_embeds is None:
+            prompt = [prompt] if isinstance(prompt, str) else prompt
+            prompt_embeds, prompt_attention_mask = self._encode_text(prompt, max_sequence_length)
+        else:
+            prompt_embeds = prompt_embeds.to(device)
+            prompt_attention_mask = prompt_attention_mask.to(device).bool()
+
+        batch_size = prompt_embeds.shape[0]
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            if negative_prompt is None:
+                negative_prompt = [None] * batch_size
+            elif isinstance(negative_prompt, str):
+                negative_prompt = [negative_prompt] * batch_size
+            negative_prompt_embeds, negative_prompt_attention_mask = self._encode_text(
+                negative_prompt, max_sequence_length
+            )
+        elif do_classifier_free_guidance:
+            negative_prompt_embeds = negative_prompt_embeds.to(device)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.to(device).bool()
+
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+        prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_images_per_prompt, dim=0)
+        if do_classifier_free_guidance:
+            negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat_interleave(
+                num_images_per_prompt, dim=0
+            )
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    def generate_vq_tokens(
+        self,
+        prompt: str | list[str],
+        height: int,
+        width: int,
+    ) -> torch.Tensor:
+        prompts = [prompt] if isinstance(prompt, str) else prompt
+        image_token_offset = 157184
+        frontend_scale = max(max(height, width) / 512, 1.0)
+        frontend_height = int(height / frontend_scale)
+        frontend_width = int(width / frontend_scale)
+        vq_height = frontend_height // 16
+        vq_width = frontend_width // 16
+        image_token_count = vq_height * vq_width
+        system_prompt = "You are a text-to-image generation assistant."
+        generated_tokens = []
+
+        for prompt in prompts:
+            text_prompt = f"<role>SYSTEM</role> {system_prompt} <role>HUMAN</role>{prompt}<role>ASSISTANT</role>"
+            text_ids = self.tokenizer(text_prompt).input_ids
+            image_info_ids = self.tokenizer(
+                f"<|image|><|reserved_token_{vq_height}|><|reserved_token_{vq_width}|><boi><|/image|>"
+            ).input_ids
+            input_ids = text_ids + image_info_ids[:-1]
+
+            uncond_prompt = (
+                f"<role>SYSTEM</role> {system_prompt} <role>HUMAN</role><uncondition><role>ASSISTANT</role>"
+            )
+            uncond_ids = self.tokenizer(uncond_prompt).input_ids + image_info_ids[:-1]
+            output_ids = self.text_encoder.generate_bd_image_logic(
+                data={
+                    "input_ids": torch.tensor(input_ids, device=self.text_encoder.device).unsqueeze(0),
+                    "uncond_ids": uncond_ids,
+                },
+                block_length=32,
+                steps=8,
+                gen_length=image_token_count,
+                cfg_scale=2.0,
+            )
+            token_ids = output_ids[0, len(input_ids) : len(input_ids) + image_token_count] - image_token_offset
+            if len(token_ids) != image_token_count:
+                raise ValueError(f"The MLLM generated {len(token_ids)} VQ tokens, expected {image_token_count}.")
+            if torch.any((token_ids < 0) | (token_ids >= self.sigvq.config.codebook_size)):
+                raise ValueError("The MLLM generated token IDs outside the SigVQ codebook.")
+            generated_tokens.append(token_ids)
+
+        return torch.stack(generated_tokens)
+
+    def check_inputs(
+        self,
+        prompt: str | list[str] | None,
+        image: PipelineImageInput | None,
+        generation_mode: str,
+        height: int,
+        width: int,
+        num_images_per_prompt: int,
+        prompt_embeds: torch.Tensor | None,
+        prompt_attention_mask: torch.Tensor | None,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_attention_mask: torch.Tensor | None,
+        callback_on_step_end_tensor_inputs: list[str],
+        num_inference_steps: int,
+    ) -> None:
+        if generation_mode not in {"text", "vq", "editing"}:
+            raise ValueError("`generation_mode` must be one of 'text', 'vq', or 'editing'.")
+        if generation_mode in {"text", "vq"} and image is not None:
+            raise ValueError(f"`image` must be omitted when `generation_mode='{generation_mode}'`.")
+        if generation_mode == "vq" and prompt is None:
+            raise ValueError("`prompt` is required when `generation_mode='vq'`.")
+        if generation_mode == "editing" and image is None:
+            raise ValueError("`image` is required when `generation_mode='editing'`.")
+        if generation_mode == "vq" and (height % 16 != 0 or width % 16 != 0):
+            raise ValueError("`height` and `width` must be divisible by 16 in VQ mode.")
+
+        required_multiple = self.latent_scale_factor * (2 if generation_mode == "editing" else 1)
+        if height <= 0 or width <= 0 or height % required_multiple != 0 or width % required_multiple != 0:
+            raise ValueError(f"`height` and `width` must be divisible by {required_multiple}.")
+        if num_inference_steps < 1:
+            raise ValueError("`num_inference_steps` must be at least 1.")
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError("Provide only one of `prompt` or `prompt_embeds`.")
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError("`prompt_attention_mask` is required with `prompt_embeds`.")
+        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
+            raise ValueError("`negative_prompt_attention_mask` is required with `negative_prompt_embeds`.")
+        if num_images_per_prompt < 1:
+            raise ValueError("`num_images_per_prompt` must be at least 1.")
+        if not all(name in self._callback_tensor_inputs for name in callback_on_step_end_tensor_inputs):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` must be chosen from {self._callback_tensor_inputs}."
+            )
+
+    def _encode_source_image(
+        self,
+        image: PipelineImageInput,
+        height: int,
+        width: int,
+        batch_size: int,
+        num_images_per_prompt: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        image = self.image_processor.preprocess(image, height=height, width=width)
+        if image.shape[0] == 1 and batch_size > 1:
+            image = image.repeat(batch_size, 1, 1, 1)
+        if image.shape[0] != batch_size:
+            raise ValueError(f"The image batch size must be 1 or {batch_size}, but is {image.shape[0]}.")
+        image = image.repeat_interleave(num_images_per_prompt, dim=0)
+
+        sigvq_pixel_values = F.interpolate(
+            image.float(),
+            size=(height // 2, width // 2),
+            mode="bilinear",
+            align_corners=False,
+        )
+        semantic_features = self.sigvq(
+            sigvq_pixel_values.to(device=self.sigvq.device, dtype=self.sigvq.dtype)
+        ).semantic_features
+
+        source_latents = self.vae.encode(image.to(device=self.vae.device, dtype=self.vae.dtype)).latent_dist.mode()
+        source_latents = self._patchify_latents(source_latents)
+        latent_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(source_latents)
+        latent_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
+            source_latents
+        )
+        source_latents = (source_latents - latent_mean) / latent_std
+        return source_latents, semantic_features
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str | list[str] | None = None,
+        image: PipelineImageInput | None = None,
+        generation_mode: str = "text",
+        negative_prompt: str | list[str] | None = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 20,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_attention_mask: torch.Tensor | None = None,
+        max_sequence_length: int = 2048,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Callable[["LLaDAImagePipeline", int, torch.Tensor, dict], dict] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+    ) -> LLaDAImagePipelineOutput | tuple:
+        r"""
+        Generate images using text-only, VQ-conditioned, or editing inference.
+
+        The timestep schedule is selected by the scheduler configuration. `use_uniform_sigmas=True` uses a uniform
+        pre-shift grid; otherwise the source Kumaraswamy schedule is used.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Text prompts that describe the generated image or requested edit.
+            image (`PipelineImageInput`, *optional*):
+                Reference image or image batch. Required in `"editing"` mode and rejected in other modes.
+            generation_mode (`str`, defaults to `"text"`):
+                Inference path. `"text"` uses only the text prompt. `"vq"` uses the MLLM to generate VQ tokens from
+                the prompt at a maximum frontend resolution of 512 before diffusion. `"editing"` uses both
+                reference-image SigVQ features and source-image latents.
+            negative_prompt (`str` or `list[str]`, *optional*):
+                Text excluded from generation. The checkpoint's empty CFG prompt is used by default.
+            height (`int`, defaults to `1024`):
+                Output image height.
+            width (`int`, defaults to `1024`):
+                Output image width.
+            num_inference_steps (`int`, defaults to `20`):
+                Number of flow-matching denoising steps.
+            guidance_scale (`float`, defaults to `4.5`):
+                Classifier-free guidance scale. Guidance is disabled at values up to `1.0`.
+            num_images_per_prompt (`int`, defaults to `1`):
+                Number of images generated per prompt.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                Random generator or generator batch used to create the initial latents.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated patchified Flux2 latents.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed, projected positive prompt embeddings.
+            prompt_attention_mask (`torch.Tensor`, *optional*):
+                Valid-token mask for `prompt_embeds`.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed, projected negative prompt embeddings.
+            negative_prompt_attention_mask (`torch.Tensor`, *optional*):
+                Valid-token mask for `negative_prompt_embeds`.
+            max_sequence_length (`int`, defaults to `2048`):
+                Maximum text sequence length before the QueryFormer tokens are appended.
+            output_type (`str`, defaults to `"pil"`):
+                Output format. Choose `"pil"`, `"np"`, `"pt"`, or `"latent"`.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return [`LLaDAImagePipelineOutput`] instead of a tuple.
+            callback_on_step_end (`Callable`, *optional*):
+                Function called after each denoising step.
+            callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`):
+                Tensor names forwarded to `callback_on_step_end`.
+
+        Returns:
+            [`LLaDAImagePipelineOutput`] or `tuple`:
+                Generated images or final patchified latents.
+        """
+        self.check_inputs(
+            prompt,
+            image,
+            generation_mode,
+            height,
+            width,
+            num_images_per_prompt,
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+            callback_on_step_end_tensor_inputs,
+            num_inference_steps,
+        )
+
+        if prompt_embeds is not None:
+            batch_size = prompt_embeds.shape[0]
+        elif isinstance(prompt, str):
+            batch_size = 1
+        else:
+            batch_size = len(prompt)
+        device = self.transformer.device
+        self._guidance_scale = guidance_scale # pylint: disable=attribute-defined-outside-init
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = (
+            self.encode_prompt(
+                prompt,
+                negative_prompt,
+                do_classifier_free_guidance,
+                num_images_per_prompt,
+                prompt_embeds,
+                prompt_attention_mask,
+                negative_prompt_embeds,
+                negative_prompt_attention_mask,
+                max_sequence_length,
+                device,
+            )
+        )
+        effective_batch_size = batch_size * num_images_per_prompt
+
+        source_latents = None
+        semantic_features = None
+        if generation_mode == "vq":
+            vq_token_ids = self.generate_vq_tokens(prompt, height, width)
+            vq_token_ids = vq_token_ids.repeat_interleave(num_images_per_prompt, dim=0)
+            semantic_features = self.sigvq(token_ids=vq_token_ids.to(self.sigvq.device)).semantic_features
+        elif generation_mode == "editing":
+            source_latents, semantic_features = self._encode_source_image(
+                image,
+                height,
+                width,
+                batch_size,
+                num_images_per_prompt,
+            )
+
+        latent_shape = (
+            effective_batch_size,
+            self.transformer.config.in_channels,
+            height // self.latent_scale_factor,
+            width // self.latent_scale_factor,
+        )
+        if latents is None:
+            latents = randn_tensor(latent_shape, generator=generator, device=device, dtype=torch.float32)
+            latents = latents.to(self.transformer.dtype).float()
+        else:
+            if latents.shape != latent_shape:
+                raise ValueError(f"Expected `latents` to have shape {latent_shape}, got {tuple(latents.shape)}.")
+            latents = latents.to(device=device, dtype=torch.float32)
+
+        if self.scheduler.config.get("use_uniform_sigmas", False):
+            # diffusers 0.39.0 does not natively support this scheduler option. Supplying the pre-shift grid
+            # explicitly preserves the behavior of the patched scheduler used by LLaDA-Image-SGLang.
+            sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1, dtype=torch.float32)[:-1].tolist()
+            self.scheduler.set_timesteps(sigmas=sigmas, device=device)
+        else:
+            schedule_steps = num_inference_steps + 1
+            schedule = torch.linspace(0.001, 1.0, schedule_steps, dtype=torch.float64)[:-1]
+            schedule = (1 - (1 - schedule**1.17) ** 0.8) ** 1.1
+            sigmas = (1 - schedule).tolist()
+            self.scheduler.set_timesteps(sigmas=sigmas, device=device)
+        timesteps = self.scheduler.timesteps
+        self._num_timesteps = len(timesteps) # pylint: disable=attribute-defined-outside-init
+
+        cond_cap_feats = [
+            embeds[mask].to(device=self.transformer.device, dtype=self.transformer.dtype)
+            for embeds, mask in zip(prompt_embeds, prompt_attention_mask.bool())
+        ]
+        if do_classifier_free_guidance:
+            uncond_cap_feats = [
+                embeds[mask].to(device=self.transformer.device, dtype=self.transformer.dtype)
+                for embeds, mask in zip(negative_prompt_embeds, negative_prompt_attention_mask.bool())
+            ]
+            cap_feats = cond_cap_feats + uncond_cap_feats
+        else:
+            cap_feats = cond_cap_feats
+
+        glm_cap_feats = None
+        source_latent_list = None
+        if semantic_features is not None:
+            cond_glm_cap_feats = [
+                features.to(device=self.transformer.device, dtype=self.transformer.dtype)
+                for features in semantic_features
+            ]
+            if source_latents is not None:
+                source_latent_list = [
+                    latent.unsqueeze(1).to(device=self.transformer.device, dtype=self.transformer.dtype)
+                    for latent in source_latents
+                ]
+            if do_classifier_free_guidance:
+                empty_glm = semantic_features.new_zeros((0, semantic_features.shape[-1])).to(
+                    device=self.transformer.device, dtype=self.transformer.dtype
+                )
+                glm_cap_feats = cond_glm_cap_feats + [empty_glm] * effective_batch_size
+                if source_latent_list is not None:
+                    source_latent_list = source_latent_list + source_latent_list
+            else:
+                glm_cap_feats = cond_glm_cap_feats
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for step_index, timestep in enumerate(timesteps):
+                latent_model_input = torch.cat([latents, latents], dim=0) if do_classifier_free_guidance else latents
+                latent_list = [latent.unsqueeze(1).to(self.transformer.dtype) for latent in latent_model_input]
+                model_timestep = (timestep / self.scheduler.config.num_train_timesteps).expand(
+                    latent_model_input.shape[0]
+                )
+
+                noise_pred = self.transformer(
+                    x=latent_list,
+                    t=model_timestep.to(self.transformer.dtype),
+                    cap_feats=cap_feats,
+                    glm_cap_feats=glm_cap_feats,
+                    source_latents=source_latent_list,
+                ).sample
+                noise_pred = -torch.stack(noise_pred, dim=0).squeeze(2).float()
+
+                if do_classifier_free_guidance:
+                    conditional_output, unconditional_output = noise_pred.chunk(2)
+                    noise_pred = unconditional_output + self.guidance_scale * (
+                        conditional_output - unconditional_output
+                    )
+
+                latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for name in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[name] = locals()[name]
+                    callback_outputs = callback_on_step_end(self, step_index, timestep, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+
+                progress_bar.update()
+
+        if output_type == "latent":
+            images = latents
+        else:
+            latents = latents.to(device=self.vae.device, dtype=self.vae.dtype)
+            latent_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents)
+            latent_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
+                latents
+            )
+            latents = latents * latent_std + latent_mean
+            latents = self._unpatchify_latents(latents)
+            images = self.vae.decode(latents, return_dict=False)[0]
+            images = self.image_processor.postprocess(images, output_type=output_type)
+
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (images,)
+        return LLaDAImagePipelineOutput(images=images)

pipelines/llada/pipeline_output.py

@@ -0,0 +1,34 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+import numpy as np
+import PIL.Image
+import torch
+
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class LLaDAImagePipelineOutput(BaseOutput):
+    """
+    Output class for the LLaDA-Image pipeline.
+
+    Args:
+        images (`list[PIL.Image.Image]`, `np.ndarray`, or `torch.Tensor`):
+            Generated images. The format is controlled by the pipeline's `output_type` argument.
+    """
+
+    images: list[PIL.Image.Image] | np.ndarray | torch.Tensor

pipelines/model_llada.py

@@ -0,0 +1,80 @@
+import diffusers
+from modules import shared, devices, sd_models, sd_hijack_te, sd_hijack_vae
+from modules.logger import log
+from pipelines import generic
+
+
+def load_llada_image(checkpoint_info, diffusers_load_config=None):
+    if diffusers_load_config is None:
+        diffusers_load_config = {}
+    repo_id = sd_models.path_to_repo(checkpoint_info)
+    sd_models.hf_auth_check(checkpoint_info)
+    log.debug(f'Load model: type=LLaDAImage repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
+
+    from pipelines.llada import LLaDAImagePipeline
+    from pipelines.llada.transformer_llada_image import LLaDAImageTransformer2DModel
+    from pipelines.llada.modeling_llada2uni_moe import LLaDA2MoeModelLM
+
+    generic.set_pipeline('LLaDAImage', LLaDAImagePipeline)
+    if repo_id is None or repo_id.lower() == 'none':
+        return None
+
+    if 'Model' in shared.opts.sdnq_quantize_weights:
+        if any(x in shared.opts.sdnq_quantize_weights_mode for x in ['2', '3', '4', '5', '6']):
+            shared.opts.sdnq_quantize_weights_mode = 'uint8'
+            log.warning('LLaDAImage: cls=LLaDAImageTransformer2DModel quant=uint8 override')
+    if 'TE' in shared.opts.sdnq_quantize_weights:
+        if any(x in shared.opts.sdnq_quantize_weights_mode_te for x in ['2', '3', '4', '5', '6']):
+            shared.opts.sdnq_quantize_weights_mode_te = 'uint8'
+            log.warning('LLaDAImage: cls=LLaDA2MoeModelLM quant=uint8 override')
+        if shared.opts.sdnq_quantize_matmul_mode_te != 'disabled':
+            shared.opts.sdnq_quantize_matmul_mode_te = 'disabled'
+            log.warning('LLaDAImage: cls=LLaDA2MoeModelLM matmul=disabled override')
+
+    transformer = generic.load_transformer(
+        repo_id,
+        cls_name=LLaDAImageTransformer2DModel,
+        load_config=diffusers_load_config,
+        modules_to_not_convert=[
+            'all_x_embedder',
+            'all_final_layer',
+            't_embedder',
+            'cap_embedder',
+            'semantic_embedder',
+            'sigvq_embedder',
+        ],
+    )
+    text_encoder = generic.load_text_encoder(
+        repo_id,
+        cls_name=LLaDA2MoeModelLM,
+        load_config=diffusers_load_config,
+        allow_shared=False,
+        trust_remote_code=True,
+        modules_to_not_convert=[
+            '.model.language_model.word_embeddings',
+            '.model.language_model.norm',
+            '.model.lm_head',
+        ],
+    )
+
+    diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['llada-image'] = LLaDAImagePipeline
+    diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['llada-image'] = LLaDAImagePipeline
+
+    pipe = LLaDAImagePipeline.from_pretrained(
+        repo_id,
+        cache_dir=shared.opts.diffusers_dir,
+        torch_dtype=devices.dtype,
+        transformer=transformer,
+        text_encoder=text_encoder,
+    )
+    pipe.task_args = {
+        'output_type': 'np',
+        'generation_mode': 'text',
+    }
+    # generation_mode = "text", "vq", "editing"
+
+    del transformer, text_encoder
+    sd_hijack_te.init_hijack(pipe)
+    sd_hijack_vae.init_hijack(pipe)
+    devices.torch_gc(force=True, reason='load')
+    return pipe