add lens img2img and inpaint

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

CHANGELOG.md

@@ -1,8 +1,8 @@
 # Change Log for SD.Next
 
-## Update for 2026-05-23
+## Update for 2026-05-25
 
-### Highlights for 2026-05-23
+### Highlights for 2026-05-25
 
 *What's New?*
 - **Anima** made it to release version, Microsoft joins the game with **Lens**
@@ -10,11 +10,11 @@
 - New **image analysis** feature and much improved **prompt enhance** capabilities which allow steering the model in real-time
 - Improved image metadata options
 
-And we have new [Contibuting** & **Development](https://vladmandic.github.io/sdnext-docs/Dev-Home/) section in docs with info on pretty much any type of development or contribution related topics - do check it out!
+And we have new [Contributing & Development](https://vladmandic.github.io/sdnext-docs/Dev-Home/) section in docs with info on pretty much any type of development or contribution related topics - do check it out!
 
-Plus continued work on modernization of codebase: UI is now fully TypeScript based and new modular LoRA loader
+Plus continued work on modernization of codebase: UI is now fully TypeScript based and we have a new modular LoRA loader
 
-### Details for 2026-05-23
+### Details for 2026-05-25
 
 - **Models**
   - [CircleStone Anima 1.0](https://huggingface.co/circlestone-labs/Anima) in *Base* and *Turbo* (distilled) variants  
@@ -23,6 +23,7 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas
     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)  
+    *note* original Lens implements only text-2-image, SD.Next adds image-2-image and inpaint workflows as well  
 - **Features**
   - **SDNQ** new quantization algorithm: *Hadamard Rotations*  
     much higher quality than base SDNQ, but runs slightly slower  
@@ -56,7 +57,7 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas
   - restore params from image metadata will now prefer *template* field if present, otherwise use *prompt* field  
     this allows to preserve original prompt in case of wildcards or styles modifying the prompt  
 - **Docs**
-  - new [Contibuting** & **Development](https://vladmandic.github.io/sdnext-docs/Dev-Home/)  
+  - new [Contributing & Development](https://vladmandic.github.io/sdnext-docs/Dev-Home/) home page  
     includes pages on *development setup, code structure, coding standards, ui development, themes, docs, hints* and more!  
 - **AI**
   - Cognitive analysis and improvements to *all* AI prompts

modules/sdnq/quantizer.py

@@ -904,7 +904,7 @@ class SDNQConfig(QuantizationConfigMixin):
         return quantization_config_dict
 
     def __str__(self):
-        return f"SDNQConfig(weights_dtype={self.weights_dtype} quantized_matmul_dtype={self.quantized_matmul_dtype} hadamard_group_size={self.hadamard_group_size} group_size={self.group_size} svd_rank={self.svd_rank} svd_steps={self.svd_steps} dynamic_loss_threshold={self.dynamic_loss_threshold} use_svd={self.use_svd} use_hadamard={self.use_hadamard} quant_conv={self.quant_conv} quant_embedding={self.quant_embedding} use_quantized_matmul={self.use_quantized_matmul} use_quantized_matmul_conv={self.use_quantized_matmul_conv} use_static_quantization={self.use_static_quantization} use_dynamic_quantization={self.use_dynamic_quantization} use_stochastic_rounding={self.use_stochastic_rounding} dequantize_fp32={self.dequantize_fp32} non_blocking={self.non_blocking} add_skip_keys={self.add_skip_keys} quantization_device={self.quantization_device} return_device={self.return_device} modules_to_not_convert={self.modules_to_not_convert} modules_to_not_use_matmul={self.modules_to_not_use_matmul} modules_dtype_dict={self.modules_dtype_dict} modules_quant_config={self.modules_quant_config} is_training={self.is_training})"
+        return f"SDNQConfig(weights_dtype={self.weights_dtype} quantization_device={self.quantization_device} return_device={self.return_device} group_size={self.group_size} use_quantized_matmul={self.use_quantized_matmul} quantized_matmul_dtype={self.quantized_matmul_dtype} quant_conv={self.quant_conv} quant_embedding={self.quant_embedding} use_quantized_matmul_conv={self.use_quantized_matmul_conv} use_static_quantization={self.use_static_quantization} use_dynamic_quantization={self.use_dynamic_quantization} dynamic_loss_threshold={self.dynamic_loss_threshold} use_stochastic_rounding={self.use_stochastic_rounding} use_hadamard={self.use_hadamard} hadamard_group_size={self.hadamard_group_size} use_svd={self.use_svd} svd_rank={self.svd_rank} svd_steps={self.svd_steps} dequantize_fp32={self.dequantize_fp32} non_blocking={self.non_blocking} add_skip_keys={self.add_skip_keys} modules_to_not_convert={self.modules_to_not_convert} modules_to_not_use_matmul={self.modules_to_not_use_matmul} modules_dtype_dict={self.modules_dtype_dict} modules_quant_config={self.modules_quant_config} )"
 
 
 import diffusers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order

pipelines/generic.py

@@ -19,7 +19,7 @@ def _loader(component):
     return 'runai' if shared.opts.runai_streamer_transformers else 'default'
 
 
-def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None):
+def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, **kwargs):
     if shared.state.interrupted:
         return None
     transformer = None
@@ -66,6 +66,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
                 cache_dir=shared.opts.hfcache_dir,
                 **load_args,
                 **quant_args,
+                **kwargs,
             )
         else:
             log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}')
@@ -82,6 +83,7 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
                 cache_dir=shared.opts.hfcache_dir,
                 **load_args,
                 **quant_args,
+                **kwargs,
             )
 
         sd_models.allow_post_quant = False # we already handled it
@@ -93,16 +95,18 @@ def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer
                 transformer.quantization_config = transformer.config.quantization_config
             elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None):
                 transformer.quantization_config = quant_args.get('quantization_config', None)
+
     except Exception as e:
         log.error(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} {e}')
         errors.display(e, 'Load')
         raise
+
     devices.torch_gc()
     shared.state.end(jobid)
     return transformer
 
 
-def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encoder", allow_quant=True, allow_shared=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None):
+def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encoder", allow_quant=True, allow_shared=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, **kwargs):
     if shared.state.interrupted:
         return None
     text_encoder = None
@@ -163,6 +167,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
                 text_encoder = nunchaku.NunchakuT5EncoderModel.from_pretrained(
                     repo_id,
                     torch_dtype=dtype,
+                    **kwargs,
                 )
                 text_encoder.quantization_method = 'SVDQuant'
             else:
@@ -179,6 +184,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
                     cache_dir=shared.opts.hfcache_dir,
                     **load_args,
                     **quant_args,
+                    **kwargs,
                 )
         elif cls_name == transformers.UMT5EncoderModel and allow_shared and shared.opts.te_shared_t5:
             if 'sdnq-uint4-svd' in repo_id.lower():
@@ -193,6 +199,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
                 subfolder=subfolder,
                 **load_args,
                 **quant_args,
+                **kwargs,
             )
         elif cls_name == transformers.Qwen2_5_VLForConditionalGeneration and allow_shared and shared.opts.te_shared_t5:
             repo_id = 'hunyuanvideo-community/HunyuanImage-2.1-Diffusers'
@@ -204,6 +211,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
                 subfolder=subfolder,
                 **load_args,
                 **quant_args,
+                **kwargs,
             )
         # Qwen3ForCausalLM - shared text encoders by hidden_size:
         # - Z-Image, Klein-4B: Qwen3-4B (hidden_size=2560)
@@ -222,6 +230,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
                 subfolder=subfolder,
                 **load_args,
                 **quant_args,
+                **kwargs,
             )
 
         # load from repo
@@ -236,6 +245,7 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
                 cache_dir=shared.opts.hfcache_dir,
                 **load_args,
                 **quant_args,
+                **kwargs,
             )
 
         sd_models.allow_post_quant = False # we already handled it
@@ -247,10 +257,12 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod
                 text_encoder.quantization_config = text_encoder.config.quantization_config
             elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None):
                 text_encoder.quantization_config = quant_args.get('quantization_config', None)
+
     except Exception as e:
         log.error(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} {e}')
         errors.display(e, 'Load')
         raise
+
     devices.torch_gc()
     shared.state.end(jobid)
     return text_encoder

pipelines/lens/__init__.py

@@ -3,6 +3,7 @@
 import diffusers as _diffusers
 import transformers as _transformers
 from .pipeline import LensPipeline, LensPipelineOutput
+from .pipeline_image import LensImg2ImgPipeline, LensInpaintPipeline
 from .reasoner import PromptReasoner
 from .resolution import RESOLUTION_BUCKETS, resolve_resolution
 from .text_encoder import LensGptOssEncoder
@@ -28,12 +29,19 @@ if not hasattr(_diffusers, "LensTransformer2DModel"):
     _diffusers.LensTransformer2DModel = LensTransformer2DModel
 if not hasattr(_diffusers, "LensPipeline"):
     _diffusers.LensPipeline = LensPipeline
+if not hasattr(_diffusers, "LensImg2ImgPipeline"):
+    _diffusers.LensImg2ImgPipeline = LensImg2ImgPipeline
+if not hasattr(_diffusers, "LensInpaintPipeline"):
+    _diffusers.LensInpaintPipeline = LensInpaintPipeline
 
+# Clean up local module references after registration.
 del _diffusers, _transformers
 
 __all__ = [
     "LensPipeline",
     "LensPipelineOutput",
+    "LensImg2ImgPipeline",
+    "LensInpaintPipeline",
     "LensTransformer2DModel",
     "LensGptOssEncoder",
     "PromptReasoner",

pipelines/lens/pipeline.py

@@ -131,8 +131,6 @@ class LensPipeline(DiffusionPipeline):
         text_encoder: LensGptOssEncoder,
         tokenizer: PreTrainedTokenizerBase,
         transformer: LensTransformer2DModel,
-        reasoner: Optional[PromptReasoner] = None,
-        use_reasoner=False,
     ) -> None:
         super().__init__()
         self.register_modules(
@@ -141,7 +139,6 @@ class LensPipeline(DiffusionPipeline):
             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
@@ -157,10 +154,9 @@ class LensPipeline(DiffusionPipeline):
                 self.transformer.config.selected_layer_index
             )
 
-        if use_reasoner and reasoner is None:
-            self.reasoner = PromptReasoner(
-                text_encoder=self.text_encoder, tokenizer=self.tokenizer
-            )
+        self.reasoner = PromptReasoner(
+            text_encoder=self.text_encoder, tokenizer=self.tokenizer
+        )
 
     # ------------------------------------------------------------------
     # Prompt encoding
@@ -295,7 +291,7 @@ class LensPipeline(DiffusionPipeline):
     def refine_prompt(
         self, prompts: Sequence[str], enable_reasoner: bool = False
     ) -> List[str]:
-        if self.reasoner is None:
+        if enable_reasoner and self.reasoner is None:
             return list(prompts)
         return self.reasoner.refine(prompts, enable=enable_reasoner)
 
@@ -370,6 +366,37 @@ class LensPipeline(DiffusionPipeline):
         latents = latents.permute(0, 1, 4, 2, 5, 3)
         return latents.reshape(b, c // 4, h * 2, w * 2)
 
+    def _pack_latents(
+        self,
+        latents: torch.Tensor,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+    ) -> torch.Tensor:
+        height = height // self.vae_scale_factor
+        width = width // self.vae_scale_factor
+        latents = latents.view(batch_size, num_channels_latents, height, 2, width, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        return latents.reshape(batch_size, height * width, num_channels_latents * 4)
+
+    def _unpack_latents(
+        self,
+        latents: torch.Tensor,
+        height: int,
+        width: int,
+        vae_scale_factor: int, # pylint: disable=unused-argument
+    ) -> torch.Tensor:
+        batch_size, _seq_len, patch_ch = latents.shape
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        return (
+            latents
+            .view(batch_size, latent_h, latent_w, patch_ch // 4, 2, 2)
+            .permute(0, 3, 1, 4, 2, 5)
+            .reshape(batch_size, patch_ch // 4, latent_h * 2, latent_w * 2)
+        )
+
     @torch.no_grad()
     def _decode(self, latents: torch.Tensor, latent_h: int, latent_w: int):
         latents = rearrange(

pipelines/lens/pipeline_image.py

@@ -0,0 +1,264 @@
+"""Lens img2img and inpainting extensions.
+
+This module adds image-to-image and inpainting support on top of the
+existing Lens text-to-image :class:`LensPipeline` without modifying the
+upstream denoising path.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from modules import devices
+from .pipeline import LensPipeline, compute_empirical_mu
+
+if TYPE_CHECKING:
+    from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+
+
+def _get_image_processor(pipe: LensPipeline) -> VaeImageProcessor:
+    processor = getattr(pipe, "lens_image_processor", None)
+    if processor is None:
+        processor = VaeImageProcessor(vae_scale_factor=pipe.vae_scale_factor // 2)
+        pipe.lens_image_processor = processor
+    return processor
+
+
+def _encode_image(
+    pipe: LensPipeline,
+    image: PipelineImageInput,
+    height: int,
+    width: int,
+    generator: Optional[torch.Generator],
+) -> torch.Tensor:
+    if isinstance(image, list):
+        image = image[0]
+
+    processor = _get_image_processor(pipe)
+    image_tensor = processor.preprocess(image, height=height, width=width)
+    image_tensor = image_tensor.to(device=devices.device, dtype=pipe.vae.dtype)
+
+    latents = pipe.vae.encode(image_tensor).latent_dist.sample(generator)
+    patched = pipe._patchify_latents(latents) # pylint: disable=protected-access
+
+    bn = pipe.vae.bn
+    mean = bn.running_mean.to(device=devices.device).view(1, -1, 1, 1)
+    std = torch.sqrt(
+        bn.running_var.to(device=devices.device).view(1, -1, 1, 1)
+        + pipe.vae.config.batch_norm_eps
+    )
+    normalized = (patched - mean) / std
+    return normalized.flatten(2).transpose(1, 2).to(devices.dtype)
+
+
+def _setup_img2img_schedule(
+    scheduler,
+    strength: float,
+    num_inference_steps: int,
+    device: torch.device,
+    mu: Optional[float] = None,
+) -> float:
+    # Avoid passing an exact zero sigma to FlowMatchEulerDiscreteScheduler,
+    # which can trigger a divide-by-zero warning during dynamic shifting.
+    min_sigma = 1e-8
+    custom_sigmas = torch.linspace(max(strength, 0.01), min_sigma, num_inference_steps).tolist()
+    scheduler.set_timesteps(sigmas=custom_sigmas, device=device, mu=mu)
+    return scheduler.sigmas[0].item()
+
+
+def _prepare_mask(
+    pipe: LensPipeline,
+    mask_image: Optional[PipelineImageInput],
+    height: int,
+    width: int,
+    device: torch.device,
+) -> torch.Tensor:
+    if isinstance(mask_image, Image.Image):
+        mask_image = mask_image.convert("L")
+    if isinstance(mask_image, Image.Image):
+        import torchvision.transforms.functional as TF
+
+        mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).to(device=device, dtype=torch.float32)
+    elif isinstance(mask_image, torch.Tensor):
+        mask_tensor = mask_image.to(device=device, dtype=torch.float32)
+        if mask_tensor.ndim == 2:
+            mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
+        elif mask_tensor.ndim == 3:
+            mask_tensor = mask_tensor.unsqueeze(0)
+    else:
+        mask_tensor = torch.ones(
+            1, 1, height, width, device=device, dtype=torch.float32
+        )
+
+    latent_h = height // pipe.vae_scale_factor
+    latent_w = width // pipe.vae_scale_factor
+    mask_latent = F.interpolate(mask_tensor, size=(latent_h, latent_w), mode="nearest")
+    mask_latent = mask_latent[:, :1, :, :]
+    return mask_latent.flatten(2).transpose(1, 2).to(device=device, dtype=pipe.vae.dtype)
+
+
+class LensImg2ImgPipeline(LensPipeline):
+    """Lens image-to-image pipeline."""
+
+    @torch.no_grad()
+    def __call__( # pylint: disable=signature-differs
+        self,
+        prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = "",
+        image: Optional[PipelineImageInput] = None,
+        strength: float = 0.8,
+        height: int = 1024,
+        width: int = 1024,
+        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,
+        negative_prompt_embeds: Optional[List[torch.Tensor]] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        enable_reasoner: bool = False,
+    ):
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        actual_sigma = _setup_img2img_schedule(
+            self.scheduler,
+            strength,
+            num_inference_steps,
+            devices.device,
+            mu=compute_empirical_mu(latent_h * latent_w, num_inference_steps),
+        )
+        init_latents = _encode_image(self, image, height, width, generator)
+        noise = randn_tensor(
+            init_latents.shape,
+            generator=generator,
+            device=devices.device,
+            dtype=devices.dtype,
+        )
+        noised = actual_sigma * noise + (1.0 - actual_sigma) * init_latents
+
+        orig_set_timesteps = self.scheduler.set_timesteps
+        self.scheduler.set_timesteps = lambda *args, **kwargs: None
+        try:
+            return super().__call__(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=num_images_per_prompt,
+                generator=generator,
+                latents=noised,
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                output_type=output_type,
+                return_dict=return_dict,
+                callback_on_step_end=callback_on_step_end,
+                callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+                max_sequence_length=max_sequence_length,
+                enable_reasoner=enable_reasoner,
+            )
+        finally:
+            self.scheduler.set_timesteps = orig_set_timesteps
+
+
+class LensInpaintPipeline(LensImg2ImgPipeline):
+    """Lens inpainting pipeline."""
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = "",
+        image: Optional[PipelineImageInput] = None,
+        mask_image: Optional[PipelineImageInput] = None,
+        strength: float = 0.8,
+        height: int = 1024,
+        width: int = 1024,
+        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,
+        negative_prompt_embeds: Optional[List[torch.Tensor]] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        enable_reasoner: bool = False,
+    ):
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        actual_sigma = _setup_img2img_schedule(
+            self.scheduler,
+            strength,
+            num_inference_steps,
+            devices.device,
+            mu=compute_empirical_mu(latent_h * latent_w, num_inference_steps),
+        )
+        init_latents = _encode_image(self, image, height, width, generator)
+        noise = randn_tensor(
+            init_latents.shape,
+            generator=generator,
+            device=devices.device,
+            dtype=devices.dtype,
+        )
+        noised = actual_sigma * noise + (1.0 - actual_sigma) * init_latents
+        mask_latent = _prepare_mask(self, mask_image, height, width, devices.device)
+
+        orig_set_timesteps = self.scheduler.set_timesteps
+        self.scheduler.set_timesteps = lambda *args, **kwargs: None
+
+        user_callback = callback_on_step_end
+
+        def blend_callback(pipe, i, t, callback_kwargs):
+            cur_latents = callback_kwargs.get("latents")
+            if cur_latents is not None:
+                sigma_next = (
+                    pipe.scheduler.sigmas[i + 1].item()
+                    if i + 1 < len(pipe.scheduler.sigmas)
+                    else 0.0
+                )
+                init_at_t = sigma_next * noise + (1.0 - sigma_next) * init_latents
+                blended = mask_latent * cur_latents + (1.0 - mask_latent) * init_at_t.to(cur_latents.dtype)
+                callback_kwargs["latents"] = blended
+            if user_callback is not None:
+                callback_kwargs = user_callback(pipe, i, t, callback_kwargs)
+            return callback_kwargs
+
+        try:
+            return LensPipeline.__call__(
+                self,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                num_images_per_prompt=num_images_per_prompt,
+                generator=generator,
+                latents=noised,
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                output_type=output_type,
+                return_dict=return_dict,
+                callback_on_step_end=blend_callback,
+                callback_on_step_end_tensor_inputs=["latents"],
+                max_sequence_length=max_sequence_length,
+                enable_reasoner=enable_reasoner,
+            )
+        finally:
+            self.scheduler.set_timesteps = orig_set_timesteps

pipelines/model_lens.py

@@ -1,3 +1,4 @@
+import diffusers
 from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
 from modules.logger import log
 from pipelines import generic
@@ -9,16 +10,14 @@ def load_lens(checkpoint_info, diffusers_load_config=None):
 
     repo_id = sd_models.path_to_repo(checkpoint_info)
     sd_models.hf_auth_check(checkpoint_info)
+    from pipelines import lens
 
     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} reasoner={shared.opts.model_lens_enable_pe} 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
+    text_encoder = generic.load_text_encoder(repo_id, cls_name=lens.LensGptOssEncoder, load_config=diffusers_load_config, allow_quant=False)
 
-    load_args['use_reasoner'] = shared.opts.model_lens_enable_pe
     pipe = lens.LensPipeline.from_pretrained(
         repo_id,
         transformer=transformer,
@@ -28,7 +27,11 @@ def load_lens(checkpoint_info, diffusers_load_config=None):
     )
     pipe.task_args = {
         "output_type": "np",
+        "enable_reasoner": shared.opts.model_lens_enable_pe,
     }
+    diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["lens"] = lens.LensPipeline
+    diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["lens"] = lens.LensImg2ImgPipeline
+    diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["lens"] = lens.LensInpaintPipeline
 
     sd_hijack_te.init_hijack(pipe)
     sd_hijack_vae.init_hijack(pipe)