add microsoft lens

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

CHANGELOG.md

@@ -1,11 +1,11 @@
 # Change Log for SD.Next
 
-## Update for 2026-05-22
+## Update for 2026-05-23
 
-### Highlights for 2026-05-22
+### Highlights for 2026-05-23
 
 *What's New?*
-- **Anima** made it to release version
+- **Anima** made it to release version, Microsoft joins the game with **Lens**
 - **SDNQ** new quantization algorithm with even higher quality
 - New **image analysis** feature and much improved **prompt enhance** capabilities which allow steering the model in real-time
 - Improved image metadata options
@@ -14,11 +14,15 @@ And we have new [Contibuting** & **Development](https://vladmandic.github.io/sdn
 
 Plus continued work on modernization of codebase: UI is now fully TypeScript based and new modular LoRA loader
 
-### Details for 2026-05-22
+### Details for 2026-05-23
 
 - **Models**
   - [CircleStone Anima 1.0](https://huggingface.co/circlestone-labs/Anima) in *Base* and *Turbo* (distilled) variants  
     in both original precision and SDNQ-4bit quantiztion  
+  - [Microsoft Lens](https://huggingface.co/microsoft/Lens) in *Standard*, *Base* and *Turbo* (distilled) variant  
+    3.8B text-to-image DiT model with 12B GPT-OSS text-encoding and Flux2 VAE  
+    oh, that 12B encoder is MoE with 3.6B activated plus its prequantized using `mxfp4`  
+    *note* Lens comes with its own prompt-refiner, enable in settings -> model options (disabled by default)  
 - **Features**
   - **SDNQ** new quantization algorithm: *Hadamard Rotations*  
     much higher quality than base SDNQ, but runs slightly slower  

data/reference-distilled.json

@@ -28,6 +28,15 @@
     "tags": "distilled",
     "extras": "sampler: Default, cfg_scale: 7.0"
   },
+  "Microsoft Lens Turbo": {
+    "path": "microsoft/Lens-Turbo",
+    "preview": "microsoft--Lens-Turbo.jpg",
+    "desc": "Microsoft Lens-Turbo is the distilled Lens variant optimized for faster text-to-image generation with fewer steps.",
+    "skip": true,
+    "tags": "distilled",
+    "size": 30.53,
+    "date": "2026 May"
+  },
   "Tencent FLUX.1 Dev SRPO": {
     "path": "vladmandic/flux.1-dev-SRPO",
     "preview": "vladmandic--flux.1-dev-SRPO.jpg",

data/reference.json

@@ -38,6 +38,22 @@
     "size": 6.94,
     "date": "2023 July"
   },
+  "Microsoft Lens": {
+    "path": "microsoft/Lens",
+    "preview": "microsoft--Lens.jpg",
+    "desc": "Microsoft Lens is a text-to-image DiT model using GPT-OSS chat-style prompt encoding and Flux2 VAE decoding.",
+    "skip": true,
+    "size": 30.53,
+    "date": "2026 May"
+  },
+  "Microsoft Lens Base": {
+    "path": "microsoft/Lens-Base",
+    "preview": "microsoft--Lens-Base.jpg",
+    "desc": "Microsoft Lens-Base is the base variant of Lens for text-to-image generation with GPT-OSS prompt features.",
+    "skip": true,
+    "size": 30.53,
+    "date": "2026 May"
+  },
   "StabilityAI Stable Cascade": {
     "path": "huggingface/stabilityai/stable-cascade",
     "skip": true,

modules/modeldata.py

@@ -94,6 +94,8 @@ def get_model_type(pipe):
         model_type = 'kolors'
     elif 'Meissonic' in name:
         model_type = 'meissonic'
+    elif 'LensPipeline' in name:
+        model_type = 'lens'
     elif 'Qwen' in name:
         model_type = 'qwen'
     elif 'ErnieImage' in name or 'ERNIE-Image' in name:

modules/sd_detect.py

@@ -99,6 +99,8 @@ def guess_by_name(fn, current_guess):
         new_guess = 'FLUX2 Klein'
     elif 'flux.2' in fn.lower():
         new_guess = 'FLUX2'
+    elif 'lens' in fn.lower():
+        new_guess = 'Lens'
     elif 'ultraflux' in fn.lower():
         new_guess = 'UltraFlux'
     elif 'flux' in fn.lower() or 'flex.1' in fn.lower():

modules/sd_models.py

@@ -486,6 +486,10 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con
             from pipelines.model_hunyuandit import load_hunyuandit
             sd_model = load_hunyuandit(checkpoint_info, diffusers_load_config)
             allow_post_quant = False
+        elif model_type in ['Lens']:
+            from pipelines.model_lens import load_lens
+            sd_model = load_lens(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Kandinsky 2.1']:
             from pipelines.model_kandinsky import load_kandinsky21
             sd_model = load_kandinsky21(checkpoint_info, diffusers_load_config)

modules/shared_items.py

@@ -54,6 +54,7 @@ pipelines = {
     'ERNIE-Image': getattr(diffusers, 'ErnieImagePipeline', None),
     'Nucleus-Image': getattr(diffusers, 'NucleusMoEImagePipeline', None),
     'Z-Image': getattr(diffusers, 'ZImagePipeline', None),
+    'Lens': getattr(diffusers, 'LensPipeline', None),
     'FLUX2': getattr(diffusers, 'Flux2Pipeline', None),
     'FLUX2 Klein': getattr(diffusers, 'Flux2KleinPipeline', None),
     'LongCat': getattr(diffusers, 'LongCatImagePipeline', None),
@@ -142,6 +143,12 @@ def get_pipelines():
             log.error(f'ONNX initialization error: {e}')
             onnx_pipelines = {}
         pipelines.update(onnx_pipelines)
+    if 'Lens' in pipelines and pipelines['Lens'] is None:
+        try:
+            import pipelines.lens as _lens
+            pipelines['Lens'] = getattr(diffusers, 'LensPipeline', None)
+        except Exception:
+            pass
     for k, v in pipelines.items():
         if k != 'Autodetect' and v is None:
             from modules.logger import log
@@ -158,6 +165,8 @@ def get_repo(model):
         return 'stabilityai/stable-diffusion-3.5-medium'
     elif model == 'FluxPipeline' or model == 'FLUX':
         return 'black-forest-labs/FLUX.1-dev'
+    elif model == 'LensPipeline' or model == 'Lens':
+        return 'microsoft/Lens'
     else:
         return None
 

modules/ui_definitions.py

@@ -113,6 +113,8 @@ def create_settings(cmd_opts):
         "model_qwen_layers": OptionInfo(2, "Qwen layered number of layers", gr.Slider, {"minimum": 2, "maximum": 9, "step": 1 }),
         "model_ernie_sep": OptionInfo("<h2>ERNIE-Image</h2>", "", gr.HTML),
         "model_ernie_enable_pe": OptionInfo(False, "Enable prompt-enhance"),
+        "model_lens_sep": OptionInfo("<h2>Lens</h2>", "", gr.HTML),
+        "model_lens_enable_pe": OptionInfo(False, "Enable prompt-enhance"),
     }))
 
     # --- Model Offloading ---

modules/vae/sd_vae_taesd.py

@@ -71,7 +71,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'}:
+    elif model_cls in {'f2', 'ernieimage', 'lens'}:
         model_cls = 'f2'
         variant = 'TAE FLUX.2'
     elif model_cls in {'sd3'}:

pipelines/lens/__init__.py

@@ -0,0 +1,42 @@
+"""Lens - minimal text-to-image inference package."""
+
+import diffusers as _diffusers
+import transformers as _transformers
+from .pipeline import LensPipeline, LensPipelineOutput
+from .reasoner import PromptReasoner
+from .resolution import RESOLUTION_BUCKETS, resolve_resolution
+from .text_encoder import LensGptOssEncoder
+from .transformer import LensTransformer2DModel
+
+# ---------------------------------------------------------------------------
+# Make our custom subclasses discoverable by ``diffusers.DiffusionPipeline``.
+#
+# When ``LensPipeline.from_pretrained`` reads ``model_index.json``, it sees
+# entries like ``["transformers", "LensGptOssEncoder"]`` and runs
+# ``getattr(importlib.import_module("transformers"), "LensGptOssEncoder")``.
+# diffusers only allow-lists the libraries ``diffusers``, ``transformers`` and
+# ``onnxruntime.training`` - any other name is interpreted as a custom .py file
+# in the repo. So we inject our subclasses into those two namespaces here.
+#
+# Importing ``lens`` is required before calling ``LensPipeline.from_pretrained``
+# (this happens automatically when the user does ``from lens import LensPipeline``).
+# ---------------------------------------------------------------------------
+
+if not hasattr(_transformers, "LensGptOssEncoder"):
+    _transformers.LensGptOssEncoder = LensGptOssEncoder
+if not hasattr(_diffusers, "LensTransformer2DModel"):
+    _diffusers.LensTransformer2DModel = LensTransformer2DModel
+if not hasattr(_diffusers, "LensPipeline"):
+    _diffusers.LensPipeline = LensPipeline
+
+del _diffusers, _transformers
+
+__all__ = [
+    "LensPipeline",
+    "LensPipelineOutput",
+    "LensTransformer2DModel",
+    "LensGptOssEncoder",
+    "PromptReasoner",
+    "RESOLUTION_BUCKETS",
+    "resolve_resolution",
+]

pipelines/lens/pipeline.py

@@ -0,0 +1,647 @@
+"""Lens text-to-image pipeline.
+
+The pipeline follows the standard ``diffusers`` component and call conventions:
+components are registered via ``register_modules`` and the call signature
+supports ``height``/``width``, ``generator``, ``prompt_embeds``, ``output_type``,
+``return_dict``, and ``callback_on_step_end``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Union
+
+import numpy as np
+import torch
+from diffusers import (
+    AutoencoderKLFlux2,
+    DiffusionPipeline,
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from einops import rearrange
+from PIL import Image
+
+from .reasoner import PromptReasoner
+from .resolution import resolve_resolution
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase
+    from .text_encoder import LensGptOssEncoder
+    from .transformer import LensTransformer2DModel
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
+    """Empirical ``mu`` for ``FlowMatchEulerDiscreteScheduler`` dynamic shift.
+
+    Constants are calibrated for the Lens inference schedule.
+    """
+    a1, b1 = 8.73809524e-05, 1.89833333
+    a2, b2 = 0.00016927, 0.45666666
+    if image_seq_len > 4300:
+        return float(a2 * image_seq_len + b2)
+    m_200 = a2 * image_seq_len + b2
+    m_10 = a1 * image_seq_len + b1
+    a = (m_200 - m_10) / 190.0
+    b = m_200 - 200.0 * a
+    return float(a * num_steps + b)
+
+
+# Chat template constants used by the Lens text encoder.
+_CHAT_SYSTEM = (
+    "Describe the image by detailing the color, shape, size, texture, "
+    "quantity, text, spatial relationships of the objects and background."
+)
+_CHAT_ASSISTANT_THINKING = "Need to generate one image according to the description."
+DEFAULT_TXT_OFFSET = 97
+
+
+# Default Lens transformer architecture.
+DEFAULT_TRANSFORMER_CONFIG = dict(
+    patch_size=2,
+    in_channels=128,
+    out_channels=32,
+    num_layers=48,
+    attention_head_dim=64,
+    num_attention_heads=24,
+    inner_dim=1536,
+    enc_hidden_dim=2880,
+    axes_dims_rope=(8, 28, 28),
+    gate_mlp=True,
+    rms_norm=True,
+    multi_layer_encoder_feature=True,
+    selected_layer_index=(5, 11, 17, 23),
+)
+
+
+@dataclass
+class LensPipelineOutput(BaseOutput):
+    """Output of :class:`LensPipeline`.
+
+    Args:
+        images: list of decoded PIL images, or a numpy array of shape
+            ``[B, H, W, C]`` when ``output_type='np'``, or the raw latent
+            tensor when ``output_type='latent'``.
+    """
+
+    images: Union[List[Image.Image], np.ndarray, torch.Tensor]
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+
+class LensPipeline(DiffusionPipeline):
+    r"""Lens text-to-image pipeline (GPT-OSS multi-layer features + Flux2 VAE).
+
+    Args:
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler used together with ``transformer`` to denoise the
+            encoded image latents.
+        vae ([`AutoencoderKLFlux2`]):
+            Flux2 VAE used to decode latents into pixel images.
+        text_encoder ([`LensGptOssEncoder`]):
+            ``GptOssForCausalLM`` subclass that exposes hidden states at the
+            configured ``selected_layer_index`` via ``encode_layers(...)``.
+        tokenizer ([`PreTrainedTokenizerBase`]):
+            GPT-OSS tokenizer.
+        transformer ([`LensTransformer2DModel`]):
+            The Lens denoising DiT.
+        reasoner ([`PromptReasoner`], *optional*):
+            Optional prompt rewriter (local OSS ``generate`` or
+            OpenAI-compatible API).
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = [
+        "latents", "prompt_embeds", "negative_prompt_embeds",
+    ]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLFlux2,
+        text_encoder: LensGptOssEncoder,
+        tokenizer: PreTrainedTokenizerBase,
+        transformer: LensTransformer2DModel,
+        reasoner: Optional[PromptReasoner] = True,
+    ) -> None:
+        super().__init__()
+        self.register_modules(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            reasoner=reasoner,
+        )
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "right"
+        # Flux2 latent tile factor (4x4 patchify) and Lens DiT in_channels=128.
+        self.vae_scale_factor = 16
+        self.latent_channels = self.transformer.config.in_channels
+        self.txt_offset = DEFAULT_TXT_OFFSET
+        self.default_sample_size = 1024
+
+        if not hasattr(self.text_encoder, "_lens_selected_layers"):
+            self.text_encoder.set_selected_layers(
+                self.transformer.config.selected_layer_index
+            )
+
+        if reasoner is not None:
+            self.reasoner = PromptReasoner(
+                text_encoder=self.text_encoder, tokenizer=self.tokenizer
+            )
+
+    # ------------------------------------------------------------------
+    # Prompt encoding
+    # ------------------------------------------------------------------
+
+    def _build_chat_inputs(
+        self, prompts: Sequence[str], max_sequence_length: int, device: torch.device
+    ):
+        rendered: List[str] = []
+        for prompt in prompts:
+            conversation = [
+                {"role": "system", "content": _CHAT_SYSTEM, "thinking": None},
+                {"role": "user", "content": prompt, "thinking": None},
+                {"role": "assistant", "thinking": _CHAT_ASSISTANT_THINKING, "content": ""},
+            ]
+            text = self.tokenizer.apply_chat_template(
+                conversation, tokenize=False, add_generation_prompt=False
+            )
+            text = text.split("<|return|>")[0]
+            rendered.append(text)
+
+        encoded = self.tokenizer(
+            rendered,
+            padding=True,
+            truncation=True,
+            max_length=max_sequence_length,
+            return_tensors="pt",
+            add_special_tokens=True,
+        )
+        return encoded["input_ids"].to(device), encoded["attention_mask"].to(device)
+
+    @torch.no_grad()
+    def _get_text_embeddings(
+        self, prompts: List[str], max_sequence_length: int, device: torch.device
+    ):
+        input_ids, attn_mask = self._build_chat_inputs(prompts, max_sequence_length, device)
+        layer_outputs = self.text_encoder.encode_layers(input_ids, attn_mask)
+
+        offset = self.txt_offset
+        if input_ids.shape[1] > offset:
+            features = [feat[:, offset:, :].contiguous() for feat in layer_outputs]
+            mask = attn_mask[:, offset:].bool()
+        else:
+            zero_shape = (input_ids.shape[0], 0, layer_outputs[0].shape[-1])
+            features = [layer_outputs[0].new_zeros(zero_shape) for _ in layer_outputs]
+            mask = torch.zeros(
+                (input_ids.shape[0], 0), dtype=torch.bool, device=device
+            )
+        return features, mask
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Union[str, List[str]] = "",
+        num_images_per_prompt: int = 1,
+        prompt_embeds: Optional[List[torch.Tensor]] = None,
+        prompt_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[List[torch.Tensor]] = None,
+        negative_prompt_mask: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 512,
+        device: Optional[torch.device] = None,
+    ):
+        """Encode positives and negatives. Returns
+        ``(prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask)``
+        where each ``*_embeds`` is a list of per-layer tensors and each
+        ``*_mask`` is a bool ``[B*N, S]`` tensor.
+
+        Each unique prompt is encoded **once**; the resulting features and mask
+        are then ``repeat_interleave``-d ``num_images_per_prompt`` times along
+        the batch axis. This preserves the ``[p0,p0,...,p1,p1,...]`` ordering
+        downstream consumers expect.
+        """
+        device = device or self._execution_device
+
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        n = int(num_images_per_prompt)
+
+        # Negatives broadcast.
+        if isinstance(negative_prompt, str):
+            negatives = [negative_prompt] * len(prompts)
+        else:
+            negatives = list(negative_prompt)
+            if len(negatives) == 1:
+                negatives = negatives * len(prompts)
+            if len(negatives) != len(prompts):
+                raise ValueError(
+                    "negative_prompt must be a string or a list of the same "
+                    "length as prompt"
+                )
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_mask = self._get_text_embeddings(
+                prompts, max_sequence_length, device
+            )
+            prompt_embeds, prompt_mask = self._repeat_for_n(prompt_embeds, prompt_mask, n)
+        elif prompt_mask is None:
+            raise ValueError("`prompt_mask` must be provided when passing `prompt_embeds`.")
+        if negative_prompt_embeds is None:
+            if all(isinstance(neg, str) and not neg.strip() for neg in negatives):
+                # Empty negatives use an unconditional branch with no text tokens.
+                negative_prompt_embeds = [
+                    feat.new_zeros(feat.shape) for feat in prompt_embeds
+                ]
+                negative_prompt_mask = torch.zeros_like(prompt_mask, dtype=torch.bool)
+            else:
+                negative_prompt_embeds, negative_prompt_mask = self._get_text_embeddings(
+                    negatives, max_sequence_length, device
+                )
+                negative_prompt_embeds, negative_prompt_mask = self._repeat_for_n(
+                    negative_prompt_embeds, negative_prompt_mask, n
+                )
+        elif negative_prompt_mask is None:
+            raise ValueError(
+                "`negative_prompt_mask` must be provided when passing "
+                "`negative_prompt_embeds`."
+            )
+        return prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask
+
+    @staticmethod
+    def _repeat_for_n(features: List[torch.Tensor], mask: torch.Tensor, n: int):
+        """Repeat each sample ``n`` times along the batch axis (interleaved)."""
+        if n == 1:
+            return features, mask
+        features = [f.repeat_interleave(n, dim=0) for f in features]
+        mask = mask.repeat_interleave(n, dim=0)
+        return features, mask
+
+    # ------------------------------------------------------------------
+    # Reasoner shim
+    # ------------------------------------------------------------------
+
+    def refine_prompt(
+        self, prompts: Sequence[str], enable_reasoner: bool = False
+    ) -> List[str]:
+        if self.reasoner is None:
+            return list(prompts)
+        print('HERE REFINE')
+        return self.reasoner.refine(prompts, enable=enable_reasoner)
+
+    # ------------------------------------------------------------------
+    # Latent prep
+    # ------------------------------------------------------------------
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        shape = (batch_size, latent_h * latent_w, num_channels_latents)
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+        return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+    # ------------------------------------------------------------------
+    # Input checks
+    # ------------------------------------------------------------------
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        prompt_embeds,
+        callback_on_step_end_tensor_inputs,
+    ) -> None:
+        if height is None or width is None:
+            raise ValueError(
+                "height and width must be provided (or use base_resolution + aspect_ratio)."
+            )
+        if height % self.vae_scale_factor or width % self.vae_scale_factor:
+            raise ValueError(
+                f"height and width must be divisible by {self.vae_scale_factor}; "
+                f"got ({height}, {width})."
+            )
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Either `prompt` or `prompt_embeds` must be provided.")
+        if callback_on_step_end_tensor_inputs is not None:
+            for k in callback_on_step_end_tensor_inputs:
+                if k not in self._callback_tensor_inputs:
+                    raise ValueError(
+                        f"callback_on_step_end_tensor_inputs entry {k!r} is not "
+                        f"in {self._callback_tensor_inputs}."
+                    )
+
+    # ------------------------------------------------------------------
+    # Decode
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _patchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = latents.shape
+        latents = latents.view(b, c, h // 2, 2, w // 2, 2)
+        latents = latents.permute(0, 1, 3, 5, 2, 4)
+        return latents.reshape(b, c * 4, h // 2, w // 2)
+
+    @staticmethod
+    def _unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = latents.shape
+        latents = latents.reshape(b, c // 4, 2, 2, h, w)
+        latents = latents.permute(0, 1, 4, 2, 5, 3)
+        return latents.reshape(b, c // 4, h * 2, w * 2)
+
+    @torch.no_grad()
+    def _decode(self, latents: torch.Tensor, latent_h: int, latent_w: int):
+        latents = rearrange(
+            latents,
+            "b (h w) (c p1 p2) -> b c (h p1) (w p2)",
+            p1=2, p2=2, h=latent_h, w=latent_w,
+        )
+        latents = latents.to(self.vae.dtype)
+        # Reverse the VAE latent normalization used by Lens. We compute the
+        # shift/scale at runtime from the live ``vae.bn`` so this stays correct
+        # under cpu-offload (where the VAE may be moved between devices).
+        bn = self.vae.bn
+        mean = bn.running_mean.view(1, -1, 1, 1)
+        var = bn.running_var.view(1, -1, 1, 1)
+        std = torch.sqrt(var + self.vae.config.batch_norm_eps)
+        shift = (-mean).to(device=latents.device, dtype=latents.dtype)
+        scale = (1.0 / std).to(device=latents.device, dtype=latents.dtype)
+        x = self._patchify_latents(latents)
+        x = x / scale - shift
+        x = self._unpatchify_latents(x)
+        return self.vae.decode(x).sample
+
+    @staticmethod
+    def _to_pil(image: torch.Tensor) -> List[Image.Image]:
+        # image: [B, C, H, W] in [-1, 1].
+        image = image.clamp(-1.0, 1.0)
+        image = (image + 1.0) * (255.0 / 2.0)
+        image = image.permute(0, 2, 3, 1).to(device="cpu", dtype=torch.uint8).numpy()
+        return [Image.fromarray(im) for im in image]
+
+    # ------------------------------------------------------------------
+    # __call__
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None, # noqa: RUF013
+        negative_prompt: Union[str, List[str]] = "",
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        base_resolution: Optional[int] = None,
+        aspect_ratio: Optional[str] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 4.0,
+        num_images_per_prompt: int = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[List[torch.Tensor]] = None,
+        prompt_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[List[torch.Tensor]] = None,
+        negative_prompt_mask: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Optional[Callable[[Any, int, int, Dict], Dict]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        enable_reasoner: bool = False,
+    ):
+        # 0. Resolution defaulting.
+        if base_resolution is not None and aspect_ratio is not None:
+            height, width = resolve_resolution(base_resolution, aspect_ratio)
+        elif height is None or width is None:
+            height = width = self.default_sample_size
+
+        # 1. Input validation.
+        self.check_inputs(
+            prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs
+        )
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        # 2. Reasoner refinement (no-op when disabled and no API).
+        if prompt is not None:
+            prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+            prompts = self.refine_prompt(prompts, enable_reasoner=enable_reasoner)
+            self._last_refined_prompts = prompts # pylint: disable=attribute-defined-outside-init
+        else:
+            prompts = None
+
+        # 3. Encode positives and negatives.
+        prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask = self.encode_prompt(
+            prompt=prompts,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_mask=prompt_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_mask=negative_prompt_mask,
+            max_sequence_length=max_sequence_length,
+            device=device,
+        )
+
+        # 4. Pad pos/neg to a shared S_txt for joint CFG batching.
+        prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask = self._align_text_features(
+            prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask
+        )
+
+        encoder_features = [
+            torch.cat([pf, nf], dim=0).to(dtype=dtype)
+            for pf, nf in zip(prompt_embeds, negative_prompt_embeds)
+        ]
+        encoder_mask = torch.cat([prompt_mask, negative_prompt_mask], dim=0)
+
+        # 5. Prepare latents.
+        batch_size = prompt_embeds[0].shape[0]
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        seq_len = latent_h * latent_w
+        latents = self.prepare_latents(
+            batch_size, self.latent_channels, height, width,
+            dtype=dtype, device=device, generator=generator, latents=latents,
+        )
+
+        # 6. Scheduler.
+        mu = compute_empirical_mu(seq_len, num_inference_steps)
+        sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)
+        self.scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu)
+
+        # 7. Denoising loop.
+        img_shapes = [(1, latent_h, latent_w)]
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(self.scheduler.timesteps):
+                timestep = t.expand(batch_size * 2).to(latents.dtype)
+                hidden_states = latents.repeat(2, 1, 1)
+
+                noise = self.transformer(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_features,
+                    encoder_hidden_states_mask=encoder_mask,
+                    timestep=timestep / 1000,
+                    img_shapes=img_shapes,
+                )
+
+                cond, uncond = noise.chunk(2)
+                comb = uncond + guidance_scale * (cond - uncond)
+                cond_norm = torch.norm(cond, dim=-1, keepdim=True)
+                comb_norm = torch.norm(comb, dim=-1, keepdim=True)
+                scale = torch.where(
+                    comb_norm > 0,
+                    cond_norm / comb_norm.clamp_min(1e-12),
+                    torch.ones_like(comb_norm),
+                )
+                noise_pred = comb * scale
+
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    cb_kwargs = {
+                        k: locals()[k] for k in callback_on_step_end_tensor_inputs
+                    }
+                    cb_out = callback_on_step_end(self, i, t, cb_kwargs)
+                    latents = cb_out.pop("latents", latents)
+                    prompt_embeds = cb_out.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = cb_out.pop(
+                        "negative_prompt_embeds", negative_prompt_embeds
+                    )
+
+                progress_bar.update()
+
+        # 8. Decode.
+        if output_type == "latent":
+            images: Any = latents
+        else:
+            decoded = self._decode(latents, latent_h, latent_w)
+            if output_type == "pil":
+                images = self._to_pil(decoded)
+            elif output_type == "np":
+                decoded = decoded.clamp(-1.0, 1.0)
+                decoded = (decoded + 1.0) * 0.5
+                images = decoded.permute(0, 2, 3, 1).to("cpu", torch.float32).numpy()
+            else:
+                raise ValueError(
+                    f"output_type must be one of 'pil', 'np', 'latent'; got {output_type!r}."
+                )
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return LensPipelineOutput(images=images)
+
+    # ------------------------------------------------------------------
+    # Misc helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _align_text_features(
+        pos_features: List[torch.Tensor],
+        pos_mask: torch.Tensor,
+        neg_features: List[torch.Tensor],
+        neg_mask: torch.Tensor,
+    ):
+        """Pad pos/neg encodings and masks to a common ``S_txt``."""
+        if not pos_features or not neg_features:
+            raise ValueError("Positive and negative text feature lists must be non-empty.")
+        if len(pos_features) != len(neg_features):
+            raise ValueError(
+                "Positive and negative text feature lists must have the same "
+                f"number of layers; got {len(pos_features)} and {len(neg_features)}."
+            )
+        seq_pos = pos_features[0].shape[1]
+        seq_neg = neg_features[0].shape[1]
+        if pos_mask.shape[1] != seq_pos:
+            raise ValueError(
+                f"prompt_mask length {pos_mask.shape[1]} does not match "
+                f"prompt feature length {seq_pos}."
+            )
+        if pos_mask.shape[0] != pos_features[0].shape[0]:
+            raise ValueError(
+                f"prompt_mask batch size {pos_mask.shape[0]} does not match "
+                f"prompt feature batch size {pos_features[0].shape[0]}."
+            )
+        if neg_mask.shape[1] != seq_neg:
+            raise ValueError(
+                f"negative_prompt_mask length {neg_mask.shape[1]} does not "
+                f"match negative prompt feature length {seq_neg}."
+            )
+        if neg_mask.shape[0] != neg_features[0].shape[0]:
+            raise ValueError(
+                f"negative_prompt_mask batch size {neg_mask.shape[0]} does "
+                f"not match negative prompt feature batch size {neg_features[0].shape[0]}."
+            )
+        if pos_features[0].shape[0] != neg_features[0].shape[0]:
+            raise ValueError(
+                "Positive and negative text features must have the same batch "
+                f"size; got {pos_features[0].shape[0]} and {neg_features[0].shape[0]}."
+            )
+        for i, feat in enumerate(pos_features):
+            if feat.shape[:2] != pos_features[0].shape[:2]:
+                raise ValueError(
+                    f"Positive feature layer {i} shape {feat.shape[:2]} does "
+                    f"not match layer 0 shape {pos_features[0].shape[:2]}."
+                )
+        for i, feat in enumerate(neg_features):
+            if feat.shape[:2] != neg_features[0].shape[:2]:
+                raise ValueError(
+                    f"Negative feature layer {i} shape {feat.shape[:2]} does "
+                    f"not match layer 0 shape {neg_features[0].shape[:2]}."
+                )
+
+        target = max(seq_pos, seq_neg)
+
+        def pad(features: List[torch.Tensor], cur: int) -> List[torch.Tensor]:
+            if cur == target:
+                return features
+            pad_len = target - cur
+            return [
+                torch.cat(
+                    [feat, feat.new_zeros((feat.shape[0], pad_len, feat.shape[-1]))],
+                    dim=1,
+                )
+                for feat in features
+            ]
+
+        def pad_mask(mask: torch.Tensor, cur: int) -> torch.Tensor:
+            if cur == target:
+                return mask
+            return torch.cat(
+                [
+                    mask,
+                    torch.zeros(
+                        (mask.shape[0], target - cur),
+                        dtype=torch.bool, device=mask.device,
+                    ),
+                ],
+                dim=1,
+            )
+
+        pos_features = pad(pos_features, seq_pos)
+        neg_features = pad(neg_features, seq_neg)
+        pos_mask = pad_mask(pos_mask.bool(), seq_pos)
+        neg_mask = pad_mask(neg_mask.bool(), seq_neg)
+        return pos_features, pos_mask, neg_features, neg_mask

pipelines/lens/reasoner.py

@@ -0,0 +1,252 @@
+"""Prompt reasoner - refines user prompts before they hit the text encoder.
+
+Decision matrix (driven by ``enable`` and whether an OpenAI-compatible API
+is configured):
+
+| ``enable``        | OpenAI API set? | Behavior                                  |
+| ----------------- | --------------- | ----------------------------------------- |
+| ``False`` (default) | no            | identity (return prompts unchanged)       |
+| ``False``         | yes             | refine via OpenAI-compatible API          |
+| ``True``          | no              | refine via the local GPT-OSS              |
+| ``True``          | yes             | refine via OpenAI-compatible API          |
+
+The OpenAI path uses any chat-completion endpoint speaking the OpenAI v1
+schema (e.g. ``vllm``, ``ollama --openai-compat``, ``together.ai``).
+"""
+
+from __future__ import annotations
+
+import re
+from typing import List, Optional, Sequence
+
+import torch
+
+
+THINK_BLOCK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
+HARMONY_FINAL_RE = re.compile(
+    r"<\|start\|>assistant(?:<\|channel\|>final)?<\|message\|>(.*?)(?:<\|return\|>|<\|end\|>|$)",
+    re.DOTALL,
+)
+HARMONY_DIRECT_FINAL_RE = re.compile(
+    r"<\|channel\|>final<\|message\|>(.*?)(?:<\|return\|>|<\|end\|>|$)",
+    re.DOTALL,
+)
+PLAIN_HARMONY_FINAL_MARKER_RE = re.compile(r"assistant\s*final\s*", re.IGNORECASE)
+PLAIN_HARMONY_DIRECT_FINAL_RE = re.compile(r"(?:^|\n)\s*final\s*", re.IGNORECASE)
+
+
+SYSTEM_PROMPT = """
+You are a prompt rewriter for a text-to-image model.
+Your task is to convert the user's input into a single, precise, descriptive image prompt suitable for a text-to-image model.
+Follow these rules strictly:
+
+1. The output must be a clear and accurate description of a single image scene, written in the style of a text-to-image prompt.
+  - Do not include explanations, reasoning, commentary, or meta text.
+  - Do not ask questions.
+  - Do not output multiple options.
+  - Do not use uncertain, speculative, or alternative wording such as "maybe", "possibly", "perhaps", "or", "might", or "could".
+
+2. Preserve the user's intended scene faithfully.
+  - Do not change the objects, entities, attributes, actions, relationships, or core setting explicitly described by the user.
+  - You may add reasonable visual details only when they help make the image concrete and coherent.
+  - Any added details must be consistent with the user's description and must not introduce new important objects or alter the meaning.
+
+3. If the image contains many main subjects of the same kind, describe each subject in detail, including humans, animals, objects, and any other prominent elements.
+  - For each subject, include its appearance, color, size, shape, material, pose, expression, and position if applicable in the scene.
+  - Make sure every main subject is clearly distinguishable from the others, such as in a scene with "4 dogs," describing each dog separately.
+
+4. The output must fully cover the scene implied by the user's input.
+  - Include the main subjects, relevant attributes, actions, spatial relationships, environment, and visible details necessary to render the scene.
+  - If the user input is already sufficiently detailed and already suitable for image generation, keep it unchanged or only make minimal edits for fluency and clarity.
+
+5. Resolve content that requires simple inference into explicit visual results when the result is unambiguous and visually representable.
+  - Example: if the user says "the answer to 2+2 is written on the blackboard", output should explicitly describe "the blackboard shows 2+2=4".
+  - Use only direct, necessary inference that is clearly implied by the user input.
+  - Do not invent hidden facts, backstory, or ambiguous details.
+
+6. Language rule:
+  - If the user input is not in English, output in the same language.
+  - Otherwise, output in English.
+
+7. Output format:
+  - Output exactly one final rewritten prompt.
+  - Do not use bullet points, numbering, JSON, XML, Markdown, or quotation marks unless they are part of the scene itself.
+
+Your goal is to produce a prompt that is concrete, visual, faithful to the user intent, and directly usable as input to a text-to-image model.
+""".strip()
+
+
+def _extract_plain_harmony_final(text: str) -> Optional[str]:
+    matches = list(PLAIN_HARMONY_FINAL_MARKER_RE.finditer(text))
+    if matches:
+        final_text = text[matches[-1].end() :].strip()
+        return final_text or None
+
+    if text.lstrip().lower().startswith("analysis"):
+        matches = list(PLAIN_HARMONY_DIRECT_FINAL_RE.finditer(text))
+        if matches:
+            final_text = text[matches[-1].end() :].strip()
+            return final_text or None
+    return None
+
+
+def _clean_reasoner_output(text: str) -> str:
+    text = text.strip()
+    final_match = None
+    for match in HARMONY_FINAL_RE.finditer(text):
+        final_match = match
+    if final_match is not None:
+        text = final_match.group(1).strip()
+    else:
+        direct_final_match = None
+        for match in HARMONY_DIRECT_FINAL_RE.finditer(text):
+            direct_final_match = match
+        if direct_final_match is not None:
+            text = direct_final_match.group(1).strip()
+        else:
+            plain_final = _extract_plain_harmony_final(text)
+            if plain_final is not None:
+                text = plain_final
+
+    text = THINK_BLOCK_RE.sub("", text).strip()
+    if "</think>" in text.lower():
+        text = re.split(r"</think>", text, flags=re.IGNORECASE)[-1].strip()
+    plain_final = _extract_plain_harmony_final(text)
+    if plain_final is not None:
+        text = plain_final
+    for token in (
+        "<|channel|>analysis<|message|>",
+        "<|start|>assistant<|channel|>analysis<|message|>",
+        "<|channel|>final<|message|>",
+        "<|start|>assistant<|channel|>final<|message|>",
+        "<|start|>assistant<|message|>",
+        "<|return|>",
+        "<|end|>",
+        "<|endoftext|>",
+        "<|im_end|>",
+    ):
+        text = text.replace(token, "")
+
+    text = text.strip()
+    if re.match(r"^(?:analysis|assistant\s*analysis)(?:\b|[A-Z])", text, flags=re.IGNORECASE | re.DOTALL):
+        return ""
+    if text.startswith("```") and text.endswith("```"):
+        lines = text.splitlines()
+        if len(lines) >= 3:
+            text = "\n".join(lines[1:-1]).strip()
+    if len(text) >= 2 and text[0] == text[-1] == '"':
+        text = text[1:-1].strip()
+    return " ".join(text.split())
+
+
+class PromptReasoner:
+    """Optional prompt rewriter, used by ``LensPipeline.refine_prompt``."""
+
+    def __init__(
+        self,
+        *,
+        text_encoder=None,
+        tokenizer=None,
+        openai_api_key: Optional[str] = None,
+        openai_base_url: Optional[str] = None,
+        openai_model: Optional[str] = None,
+        max_new_tokens: int = 4096,
+        temperature: float = 0.7,
+    ) -> None:
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.openai_api_key = openai_api_key
+        self.openai_base_url = openai_base_url
+        self.openai_model = openai_model
+        self.max_new_tokens = int(max_new_tokens)
+        self.temperature = float(temperature)
+        self._client = None  # lazily constructed
+
+    @property
+    def has_api(self) -> bool:
+        return bool(self.openai_api_key and self.openai_model)
+
+    def refine(self, prompts: Sequence[str], enable: bool) -> List[str]:
+        prompts = list(prompts)
+        # API takes precedence whenever it is configured.
+        if self.has_api:
+            return self._refine_via_api(prompts)
+        if enable:
+            if self.text_encoder is None or self.tokenizer is None:
+                raise RuntimeError(
+                    "Reasoner enabled with no API: both text_encoder and "
+                    "tokenizer must be provided to use the local GPT-OSS as "
+                    "the reasoner."
+                )
+            return self._refine_via_local(prompts)
+        return prompts
+
+    # ------------------------------------------------------------------
+    # Local GPT-OSS path
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def _refine_via_local(self, prompts: List[str]) -> List[str]:
+        refined: List[str] = []
+        for prompt in prompts:
+            system_prompt = (
+                f"{SYSTEM_PROMPT}\n\n"
+                "Keep any reasoning private. The visible answer must contain only the final rewritten prompt."
+            )
+            conversation = [
+                {"role": "system", "content": system_prompt, "thinking": None},
+                {"role": "user", "content": prompt, "thinking": None},
+            ]
+            text = self.tokenizer.apply_chat_template(
+                conversation, tokenize=False, add_generation_prompt=True, reasoning_effort="low"
+            )
+            input_ids = self.tokenizer(
+                text, return_tensors="pt", add_special_tokens=True
+            ).input_ids
+            out_ids = self.text_encoder.generate(
+                input_ids,
+                max_new_tokens=self.max_new_tokens,
+                do_sample=self.temperature > 0.0,
+                temperature=max(self.temperature, 1e-5),
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+            new_tokens = out_ids[0, input_ids.shape[1]:]
+            text_out = self.tokenizer.decode(new_tokens, skip_special_tokens=False)
+            clean_text_out = _clean_reasoner_output(text_out)
+            refined.append(clean_text_out or prompt)
+        return refined
+
+    # ------------------------------------------------------------------
+    # OpenAI-compatible API path
+    # ------------------------------------------------------------------
+
+    def _client_or_raise(self):
+        if self._client is None:
+            try:
+                from openai import OpenAI
+            except ImportError as exc:
+                raise ImportError(
+                    "openai package not installed. `pip install openai` to use "
+                    "the API-based reasoner."
+                ) from exc
+            self._client = OpenAI(
+                api_key=self.openai_api_key,
+                base_url=self.openai_base_url,
+            )
+        return self._client
+
+    def _refine_via_api(self, prompts: List[str]) -> List[str]:
+        client = self._client_or_raise()
+        out: List[str] = []
+        for prompt in prompts:
+            resp = client.chat.completions.create(
+                model=self.openai_model,
+                messages=[
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=self.max_new_tokens,
+            )
+            text = (resp.choices[0].message.content or "").strip()
+            out.append(text or prompt)
+        return out

pipelines/lens/resolution.py

@@ -0,0 +1,62 @@
+"""Resolution buckets for Lens inference.
+
+Two base resolutions (1024 and 1440) crossed with nine aspect ratios. All
+heights/widths are divisible by 16 so they tile cleanly into Flux2 VAE latents
+(downsample factor 16).
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Tuple
+
+
+# Bucket maps. Keys are aspect ratios as "W:H" strings; values are (height, width).
+# 1440 base buckets are computed as round_to_16(1024_value * 1440 / 1024).
+RESOLUTION_BUCKETS: Dict[int, Dict[str, Tuple[int, int]]] = {
+    1024: {
+        "1:2":  (1472,  736),
+        "9:16": (1376,  768),
+        "2:3":  (1248,  832),
+        "3:4":  (1152,  864),
+        "1:1":  (1024, 1024),
+        "4:3":  ( 864, 1152),
+        "3:2":  ( 832, 1248),
+        "16:9": ( 768, 1376),
+        "2:1":  ( 736, 1472),
+    },
+    1440: {
+        "1:2":  (2080, 1040),
+        "9:16": (1936, 1088),
+        "2:3":  (1760, 1168),
+        "3:4":  (1616, 1216),
+        "1:1":  (1440, 1440),
+        "4:3":  (1216, 1616),
+        "3:2":  (1168, 1760),
+        "16:9": (1088, 1936),
+        "2:1":  (1040, 2080),
+    },
+}
+
+
+SUPPORTED_BASE_RESOLUTIONS = tuple(RESOLUTION_BUCKETS.keys())
+SUPPORTED_ASPECT_RATIOS = tuple(RESOLUTION_BUCKETS[1024].keys())
+
+
+def resolve_resolution(base_resolution: int, aspect_ratio: str) -> Tuple[int, int]:
+    """Return (height, width) for the requested bucket.
+
+    Aspect ratio is interpreted as W:H (e.g. "16:9" is landscape,
+    "9:16" is portrait).
+    """
+    if base_resolution not in RESOLUTION_BUCKETS:
+        raise ValueError(
+            f"Unsupported base_resolution={base_resolution}. "
+            f"Supported: {SUPPORTED_BASE_RESOLUTIONS}"
+        )
+    table = RESOLUTION_BUCKETS[base_resolution]
+    if aspect_ratio not in table:
+        raise ValueError(
+            f"Unsupported aspect_ratio={aspect_ratio!r}. "
+            f"Supported: {SUPPORTED_ASPECT_RATIOS}"
+        )
+    return table[aspect_ratio]

pipelines/lens/text_encoder.py

@@ -0,0 +1,137 @@
+"""GPT-OSS text encoder for Lens.
+
+We subclass ``transformers.GptOssForCausalLM`` so we can:
+
+1. Return hidden states *only* at a configured layer subset (default
+   ``[5, 11, 17, 23]``), avoiding the memory cost of HF's stock
+   ``output_hidden_states=True`` which materializes every layer.
+2. Early-exit after the last selected layer, since we don't need the
+   downstream LM head at all when extracting features.
+
+Standard ``generate(...)`` is inherited unchanged and is used by the optional
+prompt reasoner.
+"""
+
+from __future__ import annotations
+
+from typing import List, Optional, Sequence
+
+import torch
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
+from transformers.models.gpt_oss.modeling_gpt_oss import GptOssForCausalLM
+
+
+class LensGptOssEncoder(GptOssForCausalLM):
+    """``GptOssForCausalLM`` subclass that exposes selected hidden states."""
+
+    def set_selected_layers(self, layer_indices: Sequence[int]) -> None:
+        layers = [int(i) for i in layer_indices]
+        if not layers:
+            raise ValueError("layer_indices must be non-empty")
+        if len(set(layers)) != len(layers):
+            raise ValueError(f"layer_indices must be unique; got {layers}")
+        if min(layers) < 0 or max(layers) >= len(self.model.layers):
+            raise ValueError(
+                f"layer_indices out of range; got {layers}, "
+                f"model has {len(self.model.layers)} layers"
+            )
+        self._lens_selected_layers = layers
+        self._lens_max_layer = max(layers)
+
+    @torch.no_grad()
+    def forward(  # type: ignore[override]
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        *args,
+        **kwargs,
+    ):
+        """Lens-specific forward.
+
+        When ``input_ids`` and ``attention_mask`` are provided AND
+        ``set_selected_layers(...)`` has been called, this returns the list of
+        hidden states at the configured selected layers (the Lens feature
+        extraction path).
+
+        Otherwise, falls back to ``GptOssForCausalLM.forward`` so that
+        ``generate(...)`` (used by the prompt reasoner) still works unchanged.
+        """
+        is_lens_feature_call = (
+            input_ids is not None
+            and attention_mask is not None
+            and hasattr(self, "_lens_selected_layers")
+            and not args
+            and not kwargs
+        )
+
+        target_device = self.model.embed_tokens.weight.device
+        if input_ids is not None and input_ids.device != target_device:
+            input_ids = input_ids.to(target_device)
+        if attention_mask is not None and attention_mask.device != target_device:
+            attention_mask = attention_mask.to(target_device)
+
+        if not is_lens_feature_call:
+            return super().forward(input_ids, attention_mask, *args, **kwargs)
+
+        model = self.model
+        inputs_embeds = model.embed_tokens(input_ids)
+        position_ids = torch.arange(
+            inputs_embeds.shape[1], device=inputs_embeds.device
+        ).unsqueeze(0).expand_as(input_ids)
+
+        mask_kwargs = {
+            "config": model.config,
+            "inputs_embeds": inputs_embeds,
+            "attention_mask": attention_mask,
+            "past_key_values": None,
+            "position_ids": position_ids,
+        }
+        causal_mask_mapping = {
+            "full_attention": create_causal_mask(**mask_kwargs),
+            "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
+        }
+
+        hidden_states = inputs_embeds
+        position_embeddings = model.rotary_emb(hidden_states, position_ids)
+
+        captured: List[torch.Tensor] = [None] * len(self._lens_selected_layers)
+        index_lookup = {idx: pos for pos, idx in enumerate(self._lens_selected_layers)}
+
+        for i, decoder_layer in enumerate(model.layers):
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask_mapping[model.config.layer_types[i]],
+                position_embeddings=position_embeddings,
+                position_ids=position_ids,
+                past_key_values=None,
+                use_cache=False,
+            )
+            if i in index_lookup:
+                captured[index_lookup[i]] = hidden_states
+            if i == self._lens_max_layer:
+                break
+
+        for pos, layer_idx in enumerate(self._lens_selected_layers):
+            if captured[pos] is None:
+                raise RuntimeError(
+                    f"Failed to capture hidden state for layer {layer_idx}"
+                )
+        return captured
+
+    def encode_layers(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.Tensor,
+    ) -> List[torch.Tensor]:
+        """Backwards-compatible alias for the Lens feature path.
+
+        Kept so existing call sites (``LensPipeline._get_text_embeddings``,
+        external users) keep working. New code should call the encoder
+        directly: ``encoder(input_ids, attention_mask)``.
+        """
+        if not hasattr(self, "_lens_selected_layers"):
+            raise RuntimeError("Call set_selected_layers(...) before encode_layers().")
+        return self(input_ids=input_ids, attention_mask=attention_mask)

pipelines/lens/transformer.py

@@ -0,0 +1,554 @@
+"""Lens denoising transformer (DiT).
+
+The model uses a double-stream architecture with joint image+text attention,
+RoPE on both streams, and SwiGLU MLPs.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward
+from diffusers.models.cache_utils import CacheMixin
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm
+
+
+# ---------------------------------------------------------------------------
+# Embeddings & RoPE
+# ---------------------------------------------------------------------------
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1.0,
+    scale: float = 1.0,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """Sinusoidal timestep embeddings (DDPM-style)."""
+    assert timesteps.ndim == 1, "Timesteps should be 1-D"
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        0, half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+    emb = torch.exp(exponent).to(timesteps.dtype)
+    emb = timesteps[:, None].float() * emb[None, :]
+    emb = scale * emb
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+    if embedding_dim % 2 == 1:
+        emb = F.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+def apply_rotary_emb_lens(
+    x: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> torch.Tensor:
+    """Apply complex-valued RoPE (Lens variant).
+
+    Args:
+        x:         [B, S, H, D] query or key tensor.
+        freqs_cis: [S, D/2] complex tensor of rotation factors.
+    """
+    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    freqs_cis = freqs_cis.unsqueeze(1)  # broadcast over heads
+    x_out = torch.view_as_real(x_complex * freqs_cis).flatten(3)
+    return x_out.type_as(x)
+
+
+class GateMLP(nn.Module):
+    """SwiGLU MLP used by the transformer blocks."""
+
+    def __init__(self, dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class LensTimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim: int) -> None:
+        super().__init__()
+        self.time_proj = Timesteps(
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000
+        )
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim
+        )
+
+    def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+        proj = self.time_proj(timestep)
+        return self.timestep_embedder(proj.to(dtype=hidden_states.dtype))
+
+
+class LensEmbedRope(nn.Module):
+    """Frame/H/W axial RoPE shared between image and text streams."""
+
+    def __init__(self, theta: int, axes_dim: List[int], scale_rope: bool = False) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        self.scale_rope = scale_rope
+        pos_index = torch.arange(4096)
+        neg_index = torch.arange(4096).flip(0) * -1 - 1
+        self.pos_freqs = torch.cat(
+            [self._rope_params(pos_index, d, theta) for d in axes_dim], dim=1
+        )
+        self.neg_freqs = torch.cat(
+            [self._rope_params(neg_index, d, theta) for d in axes_dim], dim=1
+        )
+        # Note: we deliberately do NOT register these as buffers - registering
+        # complex tensors as buffers strips the imaginary component on save/load.
+        self.rope_cache: Dict[str, torch.Tensor] = {}
+
+    @staticmethod
+    def _rope_params(index: torch.Tensor, dim: int, theta: int = 10000) -> torch.Tensor:
+        assert dim % 2 == 0
+        freqs = torch.outer(
+            index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).float().div(dim))
+        )
+        return torch.polar(torch.ones_like(freqs), freqs)
+
+    def forward(
+        self,
+        video_fhw: Union[List[Tuple[int, int, int]], Tuple[int, int, int]],
+        txt_seq_lens: Union[List[int], int],
+        device: torch.device = torch.device("cuda"),
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.pos_freqs.device != device:
+            self.pos_freqs = self.pos_freqs.to(device)
+            self.neg_freqs = self.neg_freqs.to(device)
+
+        if isinstance(video_fhw, list):
+            video_fhw = video_fhw[0]
+        if not isinstance(video_fhw, list):
+            video_fhw = [video_fhw]
+        if not isinstance(txt_seq_lens, list):
+            txt_seq_lens = [txt_seq_lens]
+        assert len(video_fhw) == 1, "video_fhw must have length 1"
+
+        vid_freqs = []
+        max_vid_index = 0
+        for idx, fhw in enumerate(video_fhw):
+            frame, height, width = fhw
+            rope_key = f"{idx}_{height}_{width}"
+            if rope_key not in self.rope_cache:
+                self.rope_cache[rope_key] = (
+                    self._compute_video_freqs(frame, height, width, idx=0).to("cpu")
+                )
+            video_freq = self.rope_cache[rope_key].to(device)
+            if self.scale_rope:
+                max_vid_index = max(height // 2, width // 2, max_vid_index)
+            else:
+                max_vid_index = max(height, width, max_vid_index)
+            vid_freqs.append(video_freq)
+
+        max_len = max(txt_seq_lens)
+        txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
+        return torch.cat(vid_freqs, dim=0), txt_freqs
+
+    def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> torch.Tensor:
+        seq_lens = frame * height * width
+        freqs_pos = self.pos_freqs.split([d // 2 for d in self.axes_dim], dim=1)
+        freqs_neg = self.neg_freqs.split([d // 2 for d in self.axes_dim], dim=1)
+
+        freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+        if self.scale_rope:
+            freqs_height = torch.cat(
+                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0
+            ).view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = torch.cat(
+                [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0
+            ).view(1, 1, width, -1).expand(frame, height, width, -1)
+        else:
+            freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
+
+        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
+        return freqs.clone().contiguous()
+
+
+# ---------------------------------------------------------------------------
+# Attention (joint image + text, plain SDPA)
+# ---------------------------------------------------------------------------
+
+
+class LensJointAttention(nn.Module):
+    """Joint image+text attention with fused QKV and SDPA backend."""
+
+    def __init__(
+        self,
+        query_dim: int,
+        added_kv_proj_dim: int,
+        dim_head: int = 64,
+        heads: int = 8,
+        out_dim: Optional[int] = None,
+        eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.heads = self.inner_dim // dim_head
+        self.dim_head = dim_head
+        self.out_dim = out_dim if out_dim is not None else query_dim
+
+        self.norm_q = RMSNorm(dim_head, eps=eps)
+        self.norm_k = RMSNorm(dim_head, eps=eps)
+        self.norm_added_q = RMSNorm(dim_head, eps=eps)
+        self.norm_added_k = RMSNorm(dim_head, eps=eps)
+
+        self.img_qkv = nn.Linear(query_dim, 3 * self.inner_dim, bias=True)
+        self.txt_qkv = nn.Linear(added_kv_proj_dim, 3 * self.inner_dim, bias=True)
+
+        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, self.out_dim, bias=True), nn.Identity()])
+        self.to_add_out = nn.Linear(self.inner_dim, query_dim, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        bsz, seq_img, _ = hidden_states.shape
+        seq_txt = encoder_hidden_states.shape[1]
+
+        # Fused QKV per stream -> split.
+        img_qkv = self.img_qkv(hidden_states).view(bsz, seq_img, 3, self.heads, self.dim_head)
+        txt_qkv = self.txt_qkv(encoder_hidden_states).view(bsz, seq_txt, 3, self.heads, self.dim_head)
+        img_q, img_k, img_v = img_qkv.unbind(dim=2)
+        txt_q, txt_k, txt_v = txt_qkv.unbind(dim=2)
+
+        # QK RMSNorm.
+        img_q = self.norm_q(img_q)
+        img_k = self.norm_k(img_k)
+        txt_q = self.norm_added_q(txt_q)
+        txt_k = self.norm_added_k(txt_k)
+
+        # RoPE.
+        img_freqs, txt_freqs = image_rotary_emb
+        if img_freqs.shape[0] < seq_img:
+            raise ValueError(
+                f"Image RoPE length {img_freqs.shape[0]} is shorter than "
+                f"image sequence length {seq_img}."
+            )
+        img_freqs = img_freqs[:seq_img]
+        img_q = apply_rotary_emb_lens(img_q, img_freqs)
+        img_k = apply_rotary_emb_lens(img_k, img_freqs)
+        if seq_txt > 0:
+            if txt_freqs.shape[0] < seq_txt:
+                raise ValueError(
+                    f"Text RoPE length {txt_freqs.shape[0]} is shorter than "
+                    f"text sequence length {seq_txt}."
+                )
+            txt_freqs = txt_freqs[:seq_txt]
+            txt_q = apply_rotary_emb_lens(txt_q, txt_freqs)
+            txt_k = apply_rotary_emb_lens(txt_k, txt_freqs)
+
+        # Joint sequence per sample, then SDPA in [B, H, S, D] layout.
+        q = torch.cat([img_q, txt_q], dim=1).transpose(1, 2)
+        k = torch.cat([img_k, txt_k], dim=1).transpose(1, 2)
+        v = torch.cat([img_v, txt_v], dim=1).transpose(1, 2)
+
+        if attention_mask is not None:
+            expected_mask_shape = (bsz, 1, 1, seq_img + seq_txt)
+            if attention_mask.shape != expected_mask_shape:
+                raise ValueError(
+                    f"attention_mask must have shape {expected_mask_shape}, "
+                    f"got {tuple(attention_mask.shape)}."
+                )
+            attention_mask = attention_mask.to(q.dtype)
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
+        out = out.transpose(1, 2).reshape(bsz, seq_img + seq_txt, -1)
+
+        img_out = self.to_out[1](self.to_out[0](out[:, :seq_img, :]))
+        txt_out = self.to_add_out(out[:, seq_img:, :])
+        return img_out, txt_out
+
+
+# ---------------------------------------------------------------------------
+# Transformer block
+# ---------------------------------------------------------------------------
+
+
+class LensTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        eps: float = 1e-6,
+        rms_norm: bool = False,
+        gate_mlp: bool = False,
+    ) -> None:
+        super().__init__()
+        self.attn = LensJointAttention(
+            query_dim=dim,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            eps=eps,
+        )
+
+        norm_cls = (lambda d: RMSNorm(d, eps=eps)) if rms_norm else (
+            lambda d: nn.LayerNorm(d, elementwise_affine=False, eps=eps)
+        )
+        if gate_mlp:
+            mlp_cls = lambda: GateMLP(dim, int(dim / 3 * 8)) # pylint: disable=unnecessary-lambda-assignment
+        else:
+            mlp_cls = lambda: FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") # pylint: disable=unnecessary-lambda-assignment
+
+        self.img_mod = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True))
+        self.img_norm1 = norm_cls(dim)
+        self.img_norm2 = norm_cls(dim)
+        self.img_mlp = mlp_cls()
+
+        self.txt_mod = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True))
+        self.txt_norm1 = norm_cls(dim)
+        self.txt_norm2 = norm_cls(dim)
+        self.txt_mlp = mlp_cls()
+
+    @staticmethod
+    def _modulate(x: torch.Tensor, mod_params: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        shift, scale, gate = mod_params.chunk(3, dim=-1)
+        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        img_mod1, img_mod2 = self.img_mod(temb).chunk(2, dim=-1)
+        txt_mod1, txt_mod2 = self.txt_mod(temb).chunk(2, dim=-1)
+
+        img_modulated, img_gate1 = self._modulate(self.img_norm1(hidden_states), img_mod1)
+        txt_modulated, txt_gate1 = self._modulate(self.txt_norm1(encoder_hidden_states), txt_mod1)
+
+        img_attn, txt_attn = self.attn(
+            hidden_states=img_modulated,
+            encoder_hidden_states=txt_modulated,
+            image_rotary_emb=image_rotary_emb,
+            attention_mask=attention_mask,
+        )
+
+        hidden_states = hidden_states + img_gate1 * img_attn
+        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn
+
+        img_modulated2, img_gate2 = self._modulate(self.img_norm2(hidden_states), img_mod2)
+        hidden_states = hidden_states + img_gate2 * self.img_mlp(img_modulated2)
+
+        txt_modulated2, txt_gate2 = self._modulate(self.txt_norm2(encoder_hidden_states), txt_mod2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2 * self.txt_mlp(txt_modulated2)
+
+        return encoder_hidden_states, hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Top-level model
+# ---------------------------------------------------------------------------
+
+
+class LensTransformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin
+):
+    """The Lens text-to-image DiT.
+
+    Supports a single conditioning stream of multi-layer text features. The
+    text features are normalized per layer, concatenated along the channel
+    axis, and projected to `inner_dim` before joining the image stream.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["LensTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+    _repeated_blocks = ["LensTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 128,
+        out_channels: Optional[int] = 32,
+        num_layers: int = 48,
+        attention_head_dim: int = 64,
+        num_attention_heads: int = 24,
+        inner_dim: int = 1536, # pylint: disable=unused-argument
+        enc_hidden_dim: int = 2880,
+        axes_dims_rope: Tuple[int, int, int] = (8, 28, 28),
+        gate_mlp: bool = True,
+        rms_norm: bool = True,
+        multi_layer_encoder_feature: bool = True,
+        selected_layer_index: Tuple[int, ...] = (5, 11, 17, 23),
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.multi_layer_encoder_feature = multi_layer_encoder_feature
+        self.selected_layer_index = list(selected_layer_index)
+
+        self.pos_embed = LensEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
+        self.time_text_embed = LensTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+        if self.multi_layer_encoder_feature:
+            self.txt_norm = nn.ModuleList(
+                [RMSNorm(enc_hidden_dim, eps=1e-5) for _ in self.selected_layer_index]
+            )
+            self.txt_in = nn.Linear(enc_hidden_dim * len(self.selected_layer_index), self.inner_dim)
+        else:
+            self.txt_norm = RMSNorm(enc_hidden_dim, eps=1e-5)
+            self.txt_in = nn.Linear(enc_hidden_dim, self.inner_dim)
+
+        self.img_in = nn.Linear(in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                LensTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    rms_norm=rms_norm,
+                    gate_mlp=gate_mlp,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
+        )
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+        encoder_hidden_states_mask: torch.Tensor,
+        timestep: torch.Tensor,
+        img_shapes: List[Tuple[int, int, int]],
+        attention_kwargs: Optional[Dict[str, Any]] = None, # pylint: disable=unused-argument
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            hidden_states:                 [B, S_img, in_channels] image latents.
+            encoder_hidden_states:         either a Tensor [B, S_txt, enc_dim]
+                                           (single-layer) or a list of such
+                                           tensors (multi-layer).
+            encoder_hidden_states_mask:    bool [B, S_txt] (True = valid).
+            timestep:                      [B] in [0, 1].
+            img_shapes:                    list with a single (frame, h_lat, w_lat).
+        """
+        bsz, img_len, _ = hidden_states.shape
+        if self.multi_layer_encoder_feature:
+            if not isinstance(encoder_hidden_states, (list, tuple)):
+                raise ValueError(
+                    "multi_layer_encoder_feature=True expects a list of "
+                    "per-layer text tensors."
+                )
+            if len(encoder_hidden_states) != len(self.selected_layer_index):
+                raise ValueError(
+                    f"Expected {len(self.selected_layer_index)} text feature "
+                    f"layers, got {len(encoder_hidden_states)}."
+                )
+            text_seq_len = encoder_hidden_states[0].shape[1]
+            for i, feat in enumerate(encoder_hidden_states):
+                if feat.shape[0] != bsz:
+                    raise ValueError(
+                        f"Text feature layer {i} batch size {feat.shape[0]} "
+                        f"does not match hidden_states batch size {bsz}."
+                    )
+                if feat.shape[1] != text_seq_len:
+                    raise ValueError(
+                        f"Text feature layer {i} sequence length {feat.shape[1]} "
+                        f"does not match layer 0 length {text_seq_len}."
+                    )
+        else:
+            if not isinstance(encoder_hidden_states, torch.Tensor):
+                raise ValueError(
+                    "multi_layer_encoder_feature=False expects a single text "
+                    "feature tensor."
+                )
+            if encoder_hidden_states.shape[0] != bsz:
+                raise ValueError(
+                    f"Text feature batch size {encoder_hidden_states.shape[0]} "
+                    f"does not match hidden_states batch size {bsz}."
+                )
+            text_seq_len = encoder_hidden_states.shape[1]
+        if encoder_hidden_states_mask.shape != (bsz, text_seq_len):
+            raise ValueError(
+                "encoder_hidden_states_mask must have shape "
+                f"{(bsz, text_seq_len)}, got {tuple(encoder_hidden_states_mask.shape)}."
+            )
+        attention_mask = self._build_joint_attention_mask(
+            encoder_hidden_states_mask, img_len
+        )
+
+        hidden_states = self.img_in(hidden_states)
+        timestep = timestep.to(hidden_states.dtype)
+
+        if self.multi_layer_encoder_feature:
+            normed = [
+                self.txt_norm[i](encoder_hidden_states[i])
+                for i in range(len(self.selected_layer_index))
+            ]
+            encoder_hidden_states = torch.cat(normed, dim=-1)
+        else:
+            encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+        temb = self.time_text_embed(timestep, hidden_states)
+
+        image_rotary_emb = self.pos_embed(
+            img_shapes, [text_seq_len], device=hidden_states.device
+        )
+
+        for block in self.transformer_blocks:
+            encoder_hidden_states, hidden_states = block(
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                attention_mask=attention_mask,
+            )
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        return self.proj_out(hidden_states)
+
+    @staticmethod
+    def _build_joint_attention_mask(
+        text_mask: torch.Tensor, img_len: int
+    ) -> torch.Tensor:
+        """Additive joint mask of shape ``[B, 1, 1, img_len + S_txt]``.
+
+        Image tokens are always valid; text positions follow ``text_mask``.
+        Padded positions hold ``-inf`` so SDPA's softmax masks them out.
+        """
+        if text_mask.dtype != torch.bool:
+            text_mask = text_mask.bool()
+        bsz = text_mask.shape[0]
+        img_ones = torch.ones(
+            (bsz, img_len), dtype=torch.bool, device=text_mask.device
+        )
+        joint = torch.cat([img_ones, text_mask], dim=1)
+        additive = torch.zeros_like(joint, dtype=torch.float32)
+        additive.masked_fill_(~joint, float("-inf"))
+        return additive[:, None, None, :]

pipelines/model_lens.py

@@ -0,0 +1,38 @@
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
+from modules.logger import log
+from pipelines import generic
+
+
+def load_lens(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)
+
+    load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+    log.debug(f'Load model: type=Lens repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+    from pipelines import lens
+
+    transformer = generic.load_transformer(repo_id, cls_name=lens.LensTransformer2DModel, load_config=diffusers_load_config)
+    text_encoder = generic.load_text_encoder(repo_id, cls_name=lens.LensGptOssEncoder, load_config=diffusers_load_config, allow_quant=False) # te is prequantized using mxfp4
+
+    if not shared.opts.model_lens_enable_pe:
+        load_args['reasoner'] = None
+
+    pipe = lens.LensPipeline.from_pretrained(
+        repo_id,
+        transformer=transformer,
+        text_encoder=text_encoder,
+        cache_dir=shared.opts.diffusers_dir,
+        **load_args,
+    )
+    pipe.task_args = {
+        "output_type": "np",
+    }
+
+    sd_hijack_te.init_hijack(pipe)
+    sd_hijack_vae.init_hijack(pipe)
+    devices.torch_gc(force=True, reason="load")
+    return pipe