Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+**/demo.png filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,192 @@
+---
+license: apache-2.0
+library_name: diffusers
+pipeline_tag: text-to-image
+base_model: CompVis/stable-diffusion-v1-4
+tags:
+  - diffusers
+  - remote-sensing
+  - optical
+  - sar
+  - infrared
+  - multimodal
+  - text-to-image
+  - lora
+  - stable-diffusion
+inference: true
+language:
+  - en
+widget:
+  - text: There is a ship in the blue water on the shore.
+    output:
+      url: demo/opt/demo.png
+  - text: There is a ship in the blue water on the shore.
+    output:
+      url: demo/sar/demo.png
+  - text: There is a ship in the blue water on the shore.
+    output:
+      url: demo/ir/demo.png
+---
+
+# BiliSakura/MMDiff-diffusers
+
+Self-contained [Diffusers](https://github.com/huggingface/diffusers) checkpoint for **MMDiff**: text-driven generation of spatially consistent **optical (OPT)**, **synthetic aperture radar (SAR)**, and **infrared (IR)** remote-sensing images from a single prompt.
+
+Converted from [`XinRan-Tang/MM-Diff`](https://huggingface.co/XinRan-Tang/MM-Diff) with the native custom pipeline in [MMDiff-diffusers](https://github.com/Bili-Sakura/MMDiff-diffusers). The optical UNet is a Stable Diffusion v1.4 backbone fine-tuned on optical image–text pairs; SAR and IR style are applied with scene LoRA adapters plus in-memory spatial feature transfer (no disk dumps under `features/`).
+
+Paper: [MMDiff: Multi-modal remote sensing image generation via cross-modality spatial feature transfer](https://doi.org/10.1016/j.isprsjprs.2026.08.018) (ISPRS Journal of Photogrammetry and Remote Sensing, 2026).
+
+## Demo
+
+Prompt: *"There is a ship in the blue water on the shore."* — scene `ship`, 256×256, 50 DDPM steps, `guidance_scale=7.5`, seed `2026`.
+
+| Optical (OPT) | SAR | Infrared (IR) |
+| --- | --- | --- |
+| ![OPT](demo/opt/demo.png) | ![SAR](demo/sar/demo.png) | ![IR](demo/ir/demo.png) |
+
+## Model details
+
+| Field | Value |
+| --- | --- |
+| Pipeline class | `MMDiffPipeline` (`pipeline.py`) |
+| Backbone | Stable Diffusion v1.4 (`UNet2DConditionModel` + `AutoencoderKL` + CLIP ViT-L/14) |
+| Scheduler | `DDPMScheduler`, 1000 training steps, `scaled_linear`, `prediction_type=epsilon` |
+| Native resolution | 256×256 |
+| Latent channels | 4 (VAE `scaling_factor=0.18215`) |
+| SAR / IR adapters | PEFT LoRA under `loras/{sar,ir}/<scene>/` |
+| Safety checker | Disabled (remote-sensing imagery) |
+| License | Apache-2.0 |
+| Training data | [`XinRan-Tang/Optical-SAR-Infrared`](https://huggingface.co/datasets/XinRan-Tang/Optical-SAR-Infrared) |
+
+### Scene LoRAs
+
+`beach`, `bridge`, `desert`, `farmland`, `lake`, `mountain`, `residential`, `river`, `ship`.
+
+SAR and IR images are decoded as single-channel (grayscale) to match the original sampling path. Optical images remain RGB.
+
+## Repo layout
+
+```text
+BiliSakura/MMDiff-diffusers/
+├── README.md
+├── pipeline.py
+├── model_index.json
+├── demo/
+│   ├── opt/demo.png
+│   ├── sar/demo.png
+│   └── ir/demo.png
+├── unet/
+├── vae/
+├── text_encoder/
+├── tokenizer/
+├── scheduler/
+└── loras/
+    ├── sar/<scene>/pytorch_lora_weights.safetensors
+    └── ir/<scene>/pytorch_lora_weights.safetensors
+```
+
+This folder is a complete Diffusers repo: load the **root**, not a nested variant.
+
+## Load from Hugging Face
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/MMDiff-diffusers",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+generator = torch.Generator(device="cpu").manual_seed(2026)
+output = pipe(
+    "There is a ship in the blue water on the shore.",
+    scene="ship",
+    height=256,
+    width=256,
+    num_inference_steps=50,
+    guidance_scale=7.5,
+    generator=generator,
+)
+output.opt[0].save("opt.png")
+output.sar[0].save("sar.png")
+output.ir[0].save("ir.png")
+```
+
+## Load from a local clone
+
+```python
+from pathlib import Path
+import torch
+from diffusers import DiffusionPipeline
+
+model_dir = Path("./MMDiff-diffusers").resolve()
+pipe = DiffusionPipeline.from_pretrained(
+    str(model_dir),
+    local_files_only=True,
+    custom_pipeline=str(model_dir / "pipeline.py"),
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+generator = torch.Generator(device="cpu").manual_seed(2026)
+output = pipe(
+    "There is a ship in the blue water on the shore.",
+    scene="ship",
+    modalities="all",  # or ["opt", "sar", "ir"]
+    height=256,
+    width=256,
+    num_inference_steps=50,
+    guidance_scale=7.5,
+    generator=generator,
+)
+output.opt[0].save("demo/opt/demo.png")
+output.sar[0].save("demo/sar/demo.png")
+output.ir[0].save("demo/ir/demo.png")
+```
+
+## Recommended inference settings
+
+| Setting | Value |
+| --- | --- |
+| Resolution | 256×256 |
+| Steps | 50 |
+| CFG (`guidance_scale`) | 7.5 |
+| `torch_dtype` | `bfloat16` |
+| Scheduler | `DDPMScheduler` (shipped) |
+| Spatial transfer | attention layers `1..9`, ResNet layer `2`, `resnet_time=1.0` |
+
+`modalities` accepts `"all"` or any subset of `"opt"`, `"sar"`, `"ir"`. SAR/IR generation requires OPT spatial features; the pipeline runs OPT first when they are not supplied via `spatial_features`.
+
+Dependencies: `diffusers`, `transformers`, `accelerate`, `peft`, `safetensors`, `torch`, `pillow`.
+
+## Interface notes
+
+- `output.images` is the first requested modality (Stable Diffusion convention); `output.opt` / `output.sar` / `output.ir` hold per-modality PIL lists.
+- `scene` selects packaged SAR/IR LoRAs. Override with `sar_lora_path` / `ir_lora_path` if needed.
+- Pass `return_spatial_features=True` to reuse captured OPT features in a later call.
+
+## Intended use
+
+Research on text-driven multi-modal remote-sensing generation, cross-modality spatial transfer, multi-modal fusion, and downstream MMRS data augmentation. Not intended as a general-purpose photorealistic image generator.
+
+## Links
+
+- Paper: [ISPRS Journal of Photogrammetry and Remote Sensing](https://www.sciencedirect.com/science/article/pii/S0924271626004089)
+- Homepage: [MMDiff](https://xinr-tang.github.io/MMDiff-homepage/)
+- Upstream weights: [`XinRan-Tang/MM-Diff`](https://huggingface.co/XinRan-Tang/MM-Diff)
+- Dataset: [`XinRan-Tang/Optical-SAR-Infrared`](https://huggingface.co/datasets/XinRan-Tang/Optical-SAR-Infrared)
+- Conversion / pipeline: [MMDiff-diffusers](https://github.com/Bili-Sakura/MMDiff-diffusers)
+
+## Citation
+
+```bibtex
+@article{tang2026mmdiff,
+  title   = {MMDiff: Multi-modal remote sensing image generation via cross-modality spatial feature transfer},
+  author  = {Tang, Haojun and Zhao, Wenda and Cui, Hengshuai and Wang, Haipeng},
+  journal = {ISPRS Journal of Photogrammetry and Remote Sensing},
+  year    = {2026},
+  doi     = {10.1016/j.isprsjprs.2026.08.018}
+}
+```

loras/ir/beach/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56a9c883d4e45f47c1c17cb5500c8ca36ae708acba816f9eb2626edda5bd6706
+size 12792952

loras/ir/bridge/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fb21c829971fc49a60eb7fd0cb6a8ae742cec3560fd1866556b362d0a7839e6
+size 12792952

loras/ir/desert/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f53c30b3ae177d5afba9e46c566b9465d7d4647ebcb08d0387c6f4c10f39738e
+size 12792952

loras/ir/farmland/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa818fecaef9f1cf6a0533e5696e450ea568563d158a9231eeaa2381d86ac716
+size 12792952

loras/ir/lake/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01cb83acbc97d43da35e137ea4bd570d19986a2486405e2ec0baa7b7b0903606
+size 12792952

loras/ir/mountain/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9830c946dcb6d60411b2b6f2cf3ba4c93e8cd740406f9a425e328b62d72c9fc0
+size 12792952

loras/ir/residential/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cff96decb990fd6214048b9a508894bd30c6450bed41db6c509d9c0f521fb5f
+size 12792952

loras/ir/river/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71662b5d960e3ebf79e16831aad00bbe9b173ce7fc9f8c3179c3f89e94465a88
+size 12792952

loras/ir/ship/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdf52809a2baf4b250a18444b0031c51af4571b65e8eaa2c808d449ffcf9ad6d
+size 12792952

loras/sar/beach/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:45e4efa97ec57352adfa818f25ee043ee2917b8921626614dfa83c4d0221e361
+size 12792952

loras/sar/bridge/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6b9bf4ddd42d62736f9b5db86201b35d9b44455a2e3af99ddf092d3727f071c
+size 12792952

loras/sar/desert/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2db5fcc87fb44db78f77961d426a7e87d206113455f58f9afd62a8af1e6c806
+size 6414992

loras/sar/farmland/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c22d077d9e76332e2c2d6188647f6dd46c3a4732544c1ce30d726086dcd626b9
+size 12792952

loras/sar/lake/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8641aad4a8718b56e2407fef16a98e13f1d5f034ce0bdf4cbb04d557f954cca7
+size 12792952

loras/sar/mountain/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f49b61257705c01ff1a6013a486e4ded12a12eccee34086022f43ab5711ff3b1
+size 12792952

loras/sar/residential/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8669346f20ecfa8477f7c22a785df591d806ff9914dcddf6c595a39a3aef00f
+size 12792952

loras/sar/river/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:146b839a06a518568619e07177336cf4030c76efdbb00f80f11f44f711867bb3
+size 12792952

loras/sar/ship/pytorch_lora_weights.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75c23a9f94465aec223ddb4289c03b8188c3e8b30696c6cacff17d596108aa69
+size 12792952

model_index.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "MMDiffPipeline",
+  "_diffusers_version": "0.32.0",
+  "feature_extractor": [
+    null,
+    null
+  ],
+  "image_encoder": [
+    null,
+    null
+  ],
+  "lora_root": "loras",
+  "requires_safety_checker": false,
+  "safety_checker": [
+    null,
+    null
+  ],
+  "scheduler": [
+    "diffusers",
+    "DDPMScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DConditionModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

pipeline.py

@@ -0,0 +1,1110 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Native Diffusers pipeline for MMDiff multi-modal remote-sensing generation."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
+
+from diffusers.image_processor import PipelineImageInput
+from diffusers.loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.models.attention_processor import Attention, AttnProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import (
+    StableDiffusionPipeline,
+    rescale_noise_cfg,
+    retrieve_timesteps,
+)
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import BaseOutput, logging, replace_example_docstring
+
+logger = logging.get_logger(__name__)
+
+DEFAULT_RESOLUTION = 256
+DEFAULT_SCENE = "ship"
+DEFAULT_ATTN_LAYERS = (1, 2, 3, 4, 5, 6, 7, 8, 9)
+DEFAULT_RESNET_LAYERS = (2,)
+SUPPORTED_MODALITIES = ("opt", "sar", "ir")
+SUPPORTED_SCENES = (
+    "beach",
+    "bridge",
+    "desert",
+    "farmland",
+    "lake",
+    "mountain",
+    "residential",
+    "river",
+    "ship",
+)
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from diffusers import DiffusionPipeline
+
+        >>> model_dir = Path("/path/to/mmdiff-diffusers")
+        >>> pipe = DiffusionPipeline.from_pretrained(
+        ...     str(model_dir),
+        ...     local_files_only=True,
+        ...     custom_pipeline=str(model_dir / "pipeline.py"),
+        ...     trust_remote_code=True,
+        ...     torch_dtype=torch.bfloat16,
+        ... )
+        >>> pipe = pipe.to("cuda")
+
+        >>> generator = torch.Generator(device="cpu").manual_seed(2026)
+        >>> output = pipe(
+        ...     "There is a ship in the blue water on the shore.",
+        ...     scene="ship",
+        ...     height=256,
+        ...     width=256,
+        ...     num_inference_steps=50,
+        ...     generator=generator,
+        ... )
+        >>> output.opt[0].save("opt.png")
+        >>> output.sar[0].save("sar.png")
+        >>> output.ir[0].save("ir.png")
+
+        >>> # Hugging Face Hub style model id: UserID/RepoID
+        >>> # Example: "XinRan-Tang/MM-Diff" after conversion, or a packaged `mmdiff-diffusers` repo.
+        ```
+"""
+
+
+def collect_up_self_attentions(unet: UNet2DConditionModel) -> list[Attention]:
+    r"""
+    Collect up-block self-attention modules in the same depth-first order used by the
+    original MMDiff hook registration.
+
+    Args:
+        unet (`UNet2DConditionModel`):
+            UNet whose `up_blocks` should be scanned.
+
+    Returns:
+        `list[Attention]`: Self-attention modules in 1-based transfer order.
+    """
+
+    modules: list[Attention] = []
+
+    def _recurse(module: torch.nn.Module) -> None:
+        if module.__class__.__name__ == "Attention":
+            if module.to_q.in_features == module.to_k.in_features:
+                modules.append(module)
+            return
+        for child in module.children():
+            _recurse(child)
+
+    _recurse(unet.up_blocks)
+    return modules
+
+
+def collect_up_resnets(unet: UNet2DConditionModel) -> list[torch.nn.Module]:
+    r"""
+    Collect up-block `ResnetBlock2D` modules in the original 0-based transfer order.
+
+    Args:
+        unet (`UNet2DConditionModel`):
+            UNet whose `up_blocks` should be scanned.
+
+    Returns:
+        `list[torch.nn.Module]`: Residual blocks in injection-index order.
+    """
+
+    modules: list[torch.nn.Module] = []
+
+    def _recurse(module: torch.nn.Module) -> None:
+        if module.__class__.__name__ == "ResnetBlock2D":
+            modules.append(module)
+            return
+        for child in module.children():
+            _recurse(child)
+
+    _recurse(unet.up_blocks)
+    return modules
+
+
+class SpatialFeatureStore:
+    r"""
+    In-memory store for OPT self-attention queries and up-block residual features.
+
+    Features are keyed by integer scheduler timestep, then by layer index. This replaces
+    the original disk dump under `features/visible_attn_maps` and `features/visible_resnet_maps`.
+    """
+
+    def __init__(self) -> None:
+        self.attn: dict[int, dict[int, torch.Tensor]] = {}
+        self.resnet: dict[int, dict[int, torch.Tensor]] = {}
+        self.current_timestep: int | None = None
+        self.mode: str = "off"
+        self.attn_layers: set[int] = set(DEFAULT_ATTN_LAYERS)
+        self.resnet_layers: set[int] = set(DEFAULT_RESNET_LAYERS)
+        self.inject_attn_timesteps: set[int] | None = None
+        self.inject_resnet_timesteps: set[int] | None = None
+
+    def reset(self) -> None:
+        r"""Clear captured features and timestep state without changing layer settings."""
+
+        self.attn = {}
+        self.resnet = {}
+        self.current_timestep = None
+
+    def set_timestep(self, timestep: int | torch.Tensor) -> None:
+        r"""
+        Record the scheduler timestep used by the current UNet forward.
+
+        Args:
+            timestep (`int` or `torch.Tensor`):
+                Scalar diffusion timestep. Tensors are stored as `int`.
+        """
+
+        self.current_timestep = int(timestep)
+
+    def _timestep_allowed(self, allowed: set[int] | None) -> bool:
+        if self.current_timestep is None:
+            return False
+        if allowed is None:
+            return True
+        return self.current_timestep in allowed
+
+    def save_attn(self, layer_idx: int, query: torch.Tensor) -> None:
+        r"""
+        Cache a self-attention query tensor for the current timestep.
+
+        Args:
+            layer_idx (`int`):
+                1-based up-block self-attention index.
+            query (`torch.Tensor`):
+                Query tensor after `head_to_batch_dim`.
+        """
+
+        if self.mode != "save" or self.current_timestep is None:
+            return
+        if layer_idx not in self.attn_layers:
+            return
+        self.attn.setdefault(self.current_timestep, {})[layer_idx] = query.detach()
+
+    def get_attn(self, layer_idx: int) -> torch.Tensor | None:
+        r"""
+        Return the cached query for the current timestep and layer, if injection is active.
+
+        Args:
+            layer_idx (`int`):
+                1-based up-block self-attention index.
+
+        Returns:
+            `torch.Tensor` or `None`: Cached query, or `None` when injection does not apply.
+        """
+
+        if self.mode != "inject" or not self._timestep_allowed(self.inject_attn_timesteps):
+            return None
+        if layer_idx not in self.attn_layers:
+            return None
+        return self.attn.get(self.current_timestep, {}).get(layer_idx)
+
+    def save_resnet(self, layer_idx: int, residual: torch.Tensor) -> None:
+        r"""
+        Cache an up-block residual tensor for the current timestep.
+
+        Args:
+            layer_idx (`int`):
+                0-based up-block ResNet index.
+            residual (`torch.Tensor`):
+                `ResnetBlock2D` output.
+        """
+
+        if self.mode != "save" or self.current_timestep is None:
+            return
+        if layer_idx not in self.resnet_layers:
+            return
+        self.resnet.setdefault(self.current_timestep, {})[layer_idx] = residual.detach()
+
+    def get_resnet(self, layer_idx: int) -> torch.Tensor | None:
+        r"""
+        Return the cached residual for the current timestep and layer, if injection is active.
+
+        Args:
+            layer_idx (`int`):
+                0-based up-block ResNet index.
+
+        Returns:
+            `torch.Tensor` or `None`: Cached residual, or `None` when injection does not apply.
+        """
+
+        if self.mode != "inject" or not self._timestep_allowed(self.inject_resnet_timesteps):
+            return None
+        if layer_idx not in self.resnet_layers:
+            return None
+        return self.resnet.get(self.current_timestep, {}).get(layer_idx)
+
+    def to_state(self) -> dict[str, Any]:
+        r"""
+        Export captured features for reuse in a later `__call__`.
+
+        Returns:
+            `dict`: Detached CPU tensors plus layer configuration.
+        """
+
+        def _cpu(store: dict[int, dict[int, torch.Tensor]]) -> dict[int, dict[int, torch.Tensor]]:
+            return {
+                timestep: {layer: tensor.detach().cpu() for layer, tensor in layers.items()}
+                for timestep, layers in store.items()
+            }
+
+        return {
+            "attn": _cpu(self.attn),
+            "resnet": _cpu(self.resnet),
+            "attn_layers": sorted(self.attn_layers),
+            "resnet_layers": sorted(self.resnet_layers),
+        }
+
+    def load_state(self, state: dict[str, Any]) -> None:
+        r"""
+        Restore features previously returned by `to_state`.
+
+        Args:
+            state (`dict`):
+                Mapping produced by `to_state`.
+        """
+
+        if not isinstance(state, dict) or "attn" not in state:
+            raise ValueError("`spatial_features` must be a dict created by MMDiffPipeline.")
+        self.attn = {int(t): {int(i): v for i, v in layers.items()} for t, layers in state["attn"].items()}
+        self.resnet = {
+            int(t): {int(i): v for i, v in layers.items()} for t, layers in state.get("resnet", {}).items()
+        }
+        if "attn_layers" in state:
+            self.attn_layers = set(int(i) for i in state["attn_layers"])
+        if "resnet_layers" in state:
+            self.resnet_layers = set(int(i) for i in state["resnet_layers"])
+
+
+class MMDiffAttnProcessor(AttnProcessor):
+    r"""
+    Attention processor that records or replaces self-attention queries during spatial transfer.
+
+    Args:
+        layer_idx (`int`):
+            1-based up-block self-attention index.
+        store (`SpatialFeatureStore`):
+            Shared feature store used by the current denoising loop.
+    """
+
+    def __init__(self, layer_idx: int, store: SpatialFeatureStore) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.store = store
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        temb: torch.Tensor | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        injected = self.store.get_attn(self.layer_idx)
+        if injected is not None:
+            query = injected.to(device=query.device, dtype=query.dtype)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        self.store.save_attn(self.layer_idx, query)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+        hidden_states = hidden_states / attn.rescale_output_factor
+        return hidden_states
+
+
+@contextmanager
+def spatial_transfer_hooks(unet: UNet2DConditionModel, store: SpatialFeatureStore):
+    r"""
+    Install native attention processors and ResNet wrappers for one denoising run.
+
+    Args:
+        unet (`UNet2DConditionModel`):
+            UNet to instrument.
+        store (`SpatialFeatureStore`):
+            Feature store read or written by the installed hooks.
+    """
+
+    attn_modules = collect_up_self_attentions(unet)
+    resnet_modules = collect_up_resnets(unet)
+    original_processors = [(module, module.processor) for module in attn_modules]
+    original_forwards = []
+
+    for layer_idx, module in enumerate(attn_modules, start=1):
+        module.set_processor(MMDiffAttnProcessor(layer_idx, store))
+
+    for layer_idx, module in enumerate(resnet_modules):
+        original_forward = module.forward
+
+        def _make_forward(orig: Callable, idx: int):
+            def wrapped(hidden_states: torch.Tensor, temb: torch.Tensor | None = None, *args: Any, **kwargs: Any):
+                output = orig(hidden_states, temb, *args, **kwargs)
+                if store.mode == "save":
+                    store.save_resnet(idx, output)
+                elif store.mode == "inject":
+                    injected = store.get_resnet(idx)
+                    if injected is not None:
+                        output = injected.to(device=output.device, dtype=output.dtype)
+                return output
+
+            return wrapped
+
+        original_forwards.append((module, original_forward))
+        module.forward = _make_forward(original_forward, layer_idx)
+
+    try:
+        yield store
+    finally:
+        for module, processor in original_processors:
+            module.set_processor(processor)
+        for module, original_forward in original_forwards:
+            module.forward = original_forward
+
+
+def normalize_modalities(modalities: str | list[str] | tuple[str, ...]) -> list[str]:
+    r"""
+    Validate and normalize the modality list passed to the pipeline.
+
+    Args:
+        modalities (`str` or sequence of `str`):
+            `"all"` or any subset of `opt`, `sar`, and `ir`.
+
+    Returns:
+        `list[str]`: Deduplicated modalities in OPT → SAR → IR order.
+    """
+
+    if isinstance(modalities, str):
+        requested = list(SUPPORTED_MODALITIES) if modalities.lower() == "all" else [modalities.lower()]
+    else:
+        requested = [str(item).lower() for item in modalities]
+
+    unknown = [item for item in requested if item not in SUPPORTED_MODALITIES]
+    if unknown:
+        raise ValueError(
+            f"Unsupported modalities {unknown}. Expected a subset of {list(SUPPORTED_MODALITIES)} or 'all'."
+        )
+    if not requested:
+        raise ValueError("At least one modality must be requested.")
+
+    ordered = [item for item in SUPPORTED_MODALITIES if item in requested]
+    return ordered
+
+
+@dataclass
+class MMDiffPipelineOutput(BaseOutput):
+    r"""
+    Output of [`MMDiffPipeline`].
+
+    Args:
+        images (`list`):
+            Images for the first requested modality, matching the Stable Diffusion `images` convention.
+        opt (`list`, *optional*):
+            Optical images when that modality was generated.
+        sar (`list`, *optional*):
+            SAR images when that modality was generated.
+        ir (`list`, *optional*):
+            Infrared images when that modality was generated.
+        nsfw_content_detected (`list[bool]`, *optional*):
+            Safety-checker flags for the primary `images` batch, if a checker is enabled.
+        spatial_features (`dict`, *optional*):
+            In-memory OPT features when `return_spatial_features=True`.
+    """
+
+    images: list[Any]
+    opt: list[Any] | None = None
+    sar: list[Any] | None = None
+    ir: list[Any] | None = None
+    nsfw_content_detected: list[bool] | None = None
+    spatial_features: dict[str, Any] | None = None
+
+
+class MMDiffPipeline(
+    StableDiffusionPipeline,
+    DiffusionPipeline,
+    TextualInversionLoaderMixin,
+    StableDiffusionLoraLoaderMixin,
+    IPAdapterMixin,
+    FromSingleFileMixin,
+):
+    r"""
+    Pipeline for jointly generating spatially consistent optical, SAR, and infrared images.
+
+    MMDiff fine-tunes a Stable Diffusion v1 UNet on optical remote-sensing pairs, then adapts
+    SAR and IR style with LoRA. During inference the OPT branch records up-block self-attention
+    queries and residual features; those features are injected into the SAR and IR branches.
+
+    Parameters:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder used to decode latents into images.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen CLIP text encoder.
+        tokenizer ([`CLIPTokenizer`]):
+            CLIP tokenizer paired with `text_encoder`.
+        unet ([`UNet2DConditionModel`]):
+            OPT-finetuned UNet. SAR/IR LoRA adapters are applied on top of this backbone.
+        scheduler ([`KarrasDiffusionSchedulers`]):
+            Denoising scheduler. The original sampling code uses [`DDPMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`], *optional*):
+            Optional safety checker. Disabled by default for remote-sensing imagery.
+        feature_extractor ([`CLIPImageProcessor`], *optional*):
+            Feature extractor used only when `safety_checker` is enabled.
+        image_encoder ([`CLIPVisionModelWithProjection`], *optional*):
+            Optional IP-Adapter image encoder.
+        requires_safety_checker (`bool`, *optional*, defaults to `False`):
+            Whether a missing safety checker should emit a warning.
+        lora_root (`str`, *optional*, defaults to `"loras"`):
+            Directory (relative to the model root) that contains `sar/<scene>` and `ir/<scene>` adapters.
+    """
+
+    model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
+    _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
+    _exclude_from_cpu_offload = ["safety_checker"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        safety_checker: StableDiffusionSafetyChecker | None = None,
+        feature_extractor: CLIPImageProcessor | None = None,
+        image_encoder: CLIPVisionModelWithProjection | None = None,
+        requires_safety_checker: bool = False,
+        lora_root: str = "loras",
+    ) -> None:
+        super().__init__(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+            image_encoder=image_encoder,
+            requires_safety_checker=requires_safety_checker,
+        )
+        self.register_to_config(lora_root=lora_root, requires_safety_checker=requires_safety_checker)
+        self.spatial_store = SpatialFeatureStore()
+        self._loaded_scene: str | None = None
+
+    def check_inputs(
+        self,
+        prompt: str | list[str] | None,
+        height: int,
+        width: int,
+        callback_steps: int | None,
+        negative_prompt: str | list[str] | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        ip_adapter_image: PipelineImageInput | None = None,
+        ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+        modalities: str | list[str] | None = None,
+        scene: str | None = None,
+        spatial_features: dict[str, Any] | None = None,
+    ) -> None:
+        r"""
+        Validate standard Stable Diffusion arguments plus MMDiff modality options.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Text prompt. Required unless `prompt_embeds` is provided.
+            height (`int`):
+                Output height in pixels. Must be divisible by the VAE scale factor.
+            width (`int`):
+                Output width in pixels. Must be divisible by the VAE scale factor.
+            callback_steps (`int`, *optional*):
+                Deprecated callback interval forwarded to the parent checker.
+            negative_prompt (`str` or `list[str]`, *optional*):
+                Negative prompt used for classifier-free guidance.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed prompt embeddings.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed negative prompt embeddings.
+            ip_adapter_image (`PipelineImageInput`, *optional*):
+                Optional IP-Adapter image.
+            ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
+                Optional precomputed IP-Adapter embeddings.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names forwarded to step-end callbacks.
+            modalities (`str` or `list[str]`, *optional*):
+                Requested modalities; validated by `normalize_modalities`.
+            scene (`str`, *optional*):
+                LoRA scene name used for SAR/IR adapters.
+            spatial_features (`dict`, *optional*):
+                Previously captured OPT features.
+        """
+
+        super().check_inputs(
+            prompt,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            callback_on_step_end_tensor_inputs,
+        )
+        if modalities is not None:
+            normalize_modalities(modalities)
+        if scene is not None and not isinstance(scene, str):
+            raise TypeError(f"`scene` must be a string, got {type(scene)}.")
+        if spatial_features is not None and not isinstance(spatial_features, dict):
+            raise TypeError("`spatial_features` must be a dict produced by this pipeline.")
+
+    def decode_latents(self, latents: torch.Tensor, single_channel: bool = False) -> torch.Tensor:
+        r"""
+        Decode latents with the VAE, optionally collapsing RGB to a single SAR/IR channel.
+
+        Args:
+            latents (`torch.Tensor`):
+                Denoised latent tensor of shape `(batch, 4, h, w)`.
+            single_channel (`bool`, *optional*, defaults to `False`):
+                If `True`, average decoded RGB channels. This matches the original
+                single-channel VAE decoder used for SAR and IR.
+
+        Returns:
+            `torch.Tensor`: Decoded images in `[-1, 1]`.
+        """
+
+        latents = latents / self.vae.config.scaling_factor
+        image = self.vae.decode(latents, return_dict=False)[0]
+        if single_channel:
+            image = image.mean(dim=1, keepdim=True)
+        return image
+
+    def resolve_lora_dir(self, modality: str, scene: str, lora_path: str | Path | None = None) -> Path:
+        r"""
+        Resolve the directory that stores a scene-specific SAR or IR LoRA.
+
+        Args:
+            modality (`str`):
+                `"sar"` or `"ir"`.
+            scene (`str`):
+                Scene name such as `"ship"` or `"beach"`.
+            lora_path (`str` or `Path`, *optional*):
+                Explicit override. When omitted, `{model_root}/{lora_root}/{modality}/{scene}` is used.
+
+        Returns:
+            `Path`: Directory expected to contain `pytorch_lora_weights.safetensors`.
+        """
+
+        if lora_path is not None:
+            return Path(lora_path)
+        root = Path(self.config.lora_root)
+        if not root.is_absolute():
+            base = getattr(self, "name_or_path", None) or "."
+            root = Path(base) / root
+        return root / modality / scene
+
+    def load_scene_loras(
+        self,
+        scene: str,
+        sar_lora_path: str | Path | None = None,
+        ir_lora_path: str | Path | None = None,
+    ) -> None:
+        r"""
+        Load SAR and IR LoRA adapters for `scene` as named PEFT adapters.
+
+        Args:
+            scene (`str`):
+                Scene used to resolve default LoRA directories.
+            sar_lora_path (`str` or `Path`, *optional*):
+                Explicit SAR adapter directory or weight file.
+            ir_lora_path (`str` or `Path`, *optional*):
+                Explicit IR adapter directory or weight file.
+        """
+
+        if self._loaded_scene == scene and sar_lora_path is None and ir_lora_path is None:
+            return
+
+        if hasattr(self, "unload_lora_weights"):
+            try:
+                self.unload_lora_weights()
+            except Exception:
+                logger.debug("No previously loaded LoRA adapters to unload.")
+
+        loaded = False
+        for modality, path in (("sar", sar_lora_path), ("ir", ir_lora_path)):
+            adapter_dir = self.resolve_lora_dir(modality, scene, path)
+            weight_file = adapter_dir if adapter_dir.is_file() else adapter_dir / "pytorch_lora_weights.safetensors"
+            if not Path(weight_file).is_file() and not adapter_dir.is_dir():
+                logger.warning("Skipping %s LoRA for scene '%s'; missing path: %s", modality, scene, adapter_dir)
+                continue
+            load_target = adapter_dir if adapter_dir.is_dir() else adapter_dir.parent
+            self.load_lora_weights(str(load_target), adapter_name=modality)
+            loaded = True
+
+        if loaded:
+            self._loaded_scene = scene
+
+    def _set_modality_adapter(self, modality: str) -> None:
+        if modality == "opt":
+            if hasattr(self, "disable_lora"):
+                try:
+                    self.disable_lora()
+                except Exception:
+                    logger.debug("LoRA disable skipped; no adapters are active.")
+            return
+        if hasattr(self, "set_adapters"):
+            try:
+                self.set_adapters(modality)
+            except Exception as error:
+                logger.warning("Could not activate the '%s' LoRA adapter: %s", modality, error)
+
+    def _postprocess_image(
+        self,
+        image: torch.Tensor,
+        output_type: str,
+        single_channel: bool,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> Any:
+        if output_type == "latent":
+            return image
+        if single_channel:
+            image, has_nsfw = self.run_safety_checker(image.repeat(1, 3, 1, 1) if image.shape[1] == 1 else image, device, dtype)
+            del has_nsfw
+            image = image.mean(dim=1, keepdim=True)
+            image = (image / 2 + 0.5).clamp(0, 1)
+            image_np = image.cpu().permute(0, 2, 3, 1).float().numpy()
+            if output_type == "np":
+                return image_np[..., 0]
+            if output_type == "pil":
+                return [
+                    self.numpy_to_pil(frame)[0].convert("L") if frame.ndim == 3 else self.numpy_to_pil(frame[..., None])[0]
+                    for frame in image_np
+                ]
+            raise ValueError(f"Unknown output_type '{output_type}'. Use 'pil', 'np', or 'latent'.")
+
+        image, has_nsfw_concept = self.run_safety_checker(image, device, dtype)
+        do_denormalize = [True] * image.shape[0] if has_nsfw_concept is None else [not flag for flag in has_nsfw_concept]
+        return self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+    def _denoise(
+        self,
+        prompt_embeds: torch.Tensor,
+        timesteps: torch.Tensor,
+        latents: torch.Tensor,
+        extra_step_kwargs: dict[str, Any],
+        timestep_cond: torch.Tensor | None,
+        added_cond_kwargs: dict[str, Any] | None,
+        store_mode: str,
+        callback: Callable | None,
+        callback_steps: int | None,
+        callback_on_step_end: Callable | None,
+        callback_on_step_end_tensor_inputs: list[str],
+        num_inference_steps: int,
+    ) -> torch.Tensor:
+        store = self.spatial_store
+        store.mode = store_mode
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        with spatial_transfer_hooks(self.unet, store), self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+                store.set_timestep(t)
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                if hasattr(self.scheduler, "scale_model_input"):
+                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep_cond=timestep_cond,
+                    cross_attention_kwargs=self.cross_attention_kwargs,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    if self.guidance_rescale > 0.0:
+                        noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
+
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {name: locals()[name] for name in callback_on_step_end_tensor_inputs if name in locals()}
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and callback_steps is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+        store.mode = "off"
+        return latents
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        timesteps: list[int] | None = None,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 7.5,
+        negative_prompt: str | list[str] | None = None,
+        num_images_per_prompt: int | None = 1,
+        eta: float = 0.0,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        ip_adapter_image: PipelineImageInput | None = None,
+        ip_adapter_image_embeds: list[torch.Tensor] | None = None,
+        output_type: str | None = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: dict[str, Any] | None = None,
+        guidance_rescale: float = 0.0,
+        clip_skip: int | None = None,
+        callback_on_step_end: Callable[..., Any] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] | None = None,
+        modalities: str | list[str] = "all",
+        scene: str = DEFAULT_SCENE,
+        sar_lora_path: str | Path | None = None,
+        ir_lora_path: str | Path | None = None,
+        attn_layers: list[int] | tuple[int, ...] | None = None,
+        resnet_layers: list[int] | tuple[int, ...] | None = None,
+        resnet_time: float = 1.0,
+        spatial_features: dict[str, Any] | None = None,
+        return_spatial_features: bool = False,
+        **kwargs: Any,
+    ) -> MMDiffPipelineOutput | tuple:
+        r"""
+        Generate optical, SAR, and/or infrared images from one text prompt.
+
+        The call follows the Stable Diffusion stage order: check inputs, define call
+        parameters, encode the prompt, prepare timesteps, prepare latents, prepare extra
+        step kwargs, then run the denoising loop. OPT is generated first so its spatial
+        features can be transferred into the SAR and IR branches.
+
+        Args:
+            prompt (`str` or `list[str]`, *optional*):
+                Text prompt that guides all requested modalities.
+            height (`int`, *optional*, defaults to `256`):
+                Output height in pixels. MMDiff was trained at 256×256.
+            width (`int`, *optional*, defaults to `256`):
+                Output width in pixels.
+            num_inference_steps (`int`, *optional*, defaults to `50`):
+                Number of denoising steps.
+            timesteps (`list[int]`, *optional*):
+                Custom descending timestep schedule.
+            sigmas (`list[float]`, *optional*):
+                Custom sigma schedule for compatible schedulers.
+            guidance_scale (`float`, *optional*, defaults to `7.5`):
+                Classifier-free guidance scale. Guidance is enabled when this value is `> 1`.
+            negative_prompt (`str` or `list[str]`, *optional*):
+                Prompt used for the unconditional branch. Defaults to empty strings.
+            num_images_per_prompt (`int`, *optional*, defaults to `1`):
+                Number of images drawn per prompt.
+            eta (`float`, *optional*, defaults to `0.0`):
+                DDIM eta. Ignored by schedulers that do not accept `eta`.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG used to sample the shared initial latents for every modality.
+            latents (`torch.Tensor`, *optional*):
+                Optional pre-sampled latents reused for every modality.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed prompt embeddings.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Precomputed unconditional embeddings.
+            ip_adapter_image (`PipelineImageInput`, *optional*):
+                Optional IP-Adapter image condition.
+            ip_adapter_image_embeds (`list[torch.Tensor]`, *optional*):
+                Optional precomputed IP-Adapter embeddings.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                `"pil"`, `"np"`, or `"latent"`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return [`MMDiffPipelineOutput`].
+            cross_attention_kwargs (`dict`, *optional*):
+                Extra kwargs forwarded to attention processors.
+            guidance_rescale (`float`, *optional*, defaults to `0.0`):
+                Optional guidance rescale factor.
+            clip_skip (`int`, *optional*):
+                Number of CLIP layers to skip when encoding prompts.
+            callback_on_step_end (`Callable`, *optional*):
+                Optional per-step callback.
+            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
+                Tensor names passed to `callback_on_step_end`.
+            modalities (`str` or `list[str]`, *optional*, defaults to `"all"`):
+                `"all"` or any subset of `"opt"`, `"sar"`, `"ir"`.
+            scene (`str`, *optional*, defaults to `"ship"`):
+                Scene used to resolve packaged SAR/IR LoRA adapters.
+            sar_lora_path (`str` or `Path`, *optional*):
+                Override for the SAR LoRA directory or weight file.
+            ir_lora_path (`str` or `Path`, *optional*):
+                Override for the IR LoRA directory or weight file.
+            attn_layers (`list[int]`, *optional*):
+                1-based up-block self-attention layers to transfer. Defaults to `1..9`.
+            resnet_layers (`list[int]`, *optional*):
+                0-based up-block ResNet layers to transfer. Defaults to `(2,)`.
+            resnet_time (`float`, *optional*, defaults to `1.0`):
+                Fraction of the early timestep schedule that receives ResNet injection.
+            spatial_features (`dict`, *optional*):
+                Features from a previous OPT run. When omitted, OPT is run first whenever
+                SAR or IR generation needs transfer features.
+            return_spatial_features (`bool`, *optional*, defaults to `False`):
+                If `True`, include the captured OPT features in the output.
+
+        Examples:
+
+        Returns:
+            [`MMDiffPipelineOutput`] or `tuple`:
+                Generated images grouped by modality. `images` is the first requested modality.
+        """
+
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs or ["latents"]
+        requested = normalize_modalities(modalities)
+        attn_layers = tuple(DEFAULT_ATTN_LAYERS if attn_layers is None else attn_layers)
+        resnet_layers = tuple(DEFAULT_RESNET_LAYERS if resnet_layers is None else resnet_layers)
+
+        height = DEFAULT_RESOLUTION if height is None else height
+        width = DEFAULT_RESOLUTION if width is None else width
+
+        # 1. Check inputs
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            callback_on_step_end_tensor_inputs,
+            modalities=requested,
+            scene=scene,
+            spatial_features=spatial_features,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        num_images_per_prompt = 1 if num_images_per_prompt is None else num_images_per_prompt
+        device = self._execution_device
+
+        needs_transfer = any(modality in requested for modality in ("sar", "ir"))
+        run_opt = "opt" in requested or (needs_transfer and spatial_features is None)
+        if needs_transfer and not run_opt and spatial_features is None:
+            raise ValueError("SAR/IR generation requires OPT spatial features. Run OPT first or pass `spatial_features`.")
+
+        if any(modality in requested for modality in ("sar", "ir")):
+            self.load_scene_loras(scene, sar_lora_path=sar_lora_path, ir_lora_path=ir_lora_path)
+
+        self.spatial_store.reset()
+        self.spatial_store.attn_layers = set(attn_layers)
+        self.spatial_store.resnet_layers = set(resnet_layers)
+        if spatial_features is not None:
+            self.spatial_store.load_state(spatial_features)
+
+        # 3. Encode input condition
+        lora_scale = self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            self.do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+        )
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        added_cond_kwargs = None
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+            added_cond_kwargs = {"image_embeds": image_embeds}
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas
+        )
+        timestep_values = [int(step) for step in timesteps]
+        self.spatial_store.inject_attn_timesteps = set(timestep_values)
+        cutoff = max(1, int(len(timestep_values) * resnet_time)) if resnet_time > 0 else 0
+        self.spatial_store.inject_resnet_timesteps = set(timestep_values[:cutoff])
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        init_latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        timestep_cond = None
+        if getattr(self.unet.config, "time_cond_proj_dim", None) is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=init_latents.dtype)
+
+        # 7. Run denoising loop for each requested branch
+        generated: dict[str, Any] = {}
+        primary_nsfw = None
+        for modality in (("opt",) if run_opt else ()) + tuple(item for item in requested if item != "opt"):
+            self._set_modality_adapter(modality)
+            store_mode = "save" if modality == "opt" else "inject"
+            latents_in = init_latents.clone()
+            latents_out = self._denoise(
+                prompt_embeds=prompt_embeds,
+                timesteps=timesteps,
+                latents=latents_in,
+                extra_step_kwargs=extra_step_kwargs,
+                timestep_cond=timestep_cond,
+                added_cond_kwargs=added_cond_kwargs,
+                store_mode=store_mode,
+                callback=callback,
+                callback_steps=callback_steps,
+                callback_on_step_end=callback_on_step_end,
+                callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+                num_inference_steps=num_inference_steps,
+            )
+            if modality not in requested:
+                continue
+            if output_type == "latent":
+                images = latents_out
+                has_nsfw = None
+            else:
+                single_channel = modality in {"sar", "ir"}
+                decoded = self.decode_latents(latents_out, single_channel=single_channel)
+                if single_channel:
+                    images = self._postprocess_image(
+                        decoded, output_type=output_type, single_channel=True, dtype=prompt_embeds.dtype, device=device
+                    )
+                    has_nsfw = None
+                else:
+                    images, has_nsfw = self.run_safety_checker(decoded, device, prompt_embeds.dtype)
+                    do_denormalize = [True] * images.shape[0] if has_nsfw is None else [not flag for flag in has_nsfw]
+                    images = self.image_processor.postprocess(
+                        images, output_type=output_type, do_denormalize=do_denormalize
+                    )
+            generated[modality] = images
+            if primary_nsfw is None:
+                primary_nsfw = has_nsfw
+
+        self.maybe_free_model_hooks()
+        feature_state = self.spatial_store.to_state() if return_spatial_features else None
+        images = generated.get(requested[0])
+        if not return_dict:
+            return (images, primary_nsfw)
+
+        return MMDiffPipelineOutput(
+            images=images,
+            opt=generated.get("opt"),
+            sar=generated.get("sar"),
+            ir=generated.get("ir"),
+            nsfw_content_detected=primary_nsfw,
+            spatial_features=feature_state,
+        )
+
+
+__all__ = ["MMDiffPipeline", "MMDiffPipelineOutput", "SpatialFeatureStore"]

scheduler/scheduler_config.json

@@ -0,0 +1,15 @@
+{
+  "_class_name": "DDPMScheduler",
+  "_diffusers_version": "0.32.0.dev0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "set_alpha_to_one": false,
+  "skip_prk_steps": true,
+  "steps_offset": 1,
+  "timestep_spacing": "leading",
+  "trained_betas": null
+}

text_encoder/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "/media/ubuntun/hdd/THj/model/stable-diffusion-v1-4",
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 512,
+  "torch_dtype": "float16",
+  "transformers_version": "4.46.1",
+  "vocab_size": 49408
+}

text_encoder/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:660c6f5b1abae9dc498ac2d21e1347d2abdb0cf6c0c0c8576cd796491d9a6cdd
+size 246144152

tokenizer/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

unet/config.json

@@ -0,0 +1,68 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.32.0.dev0",
+  "_name_or_path": "/media/ubuntun/hdd/THj/model/stable-diffusion-v1-4",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 768,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

unet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6563affab036a21949693582693155259d70e390ec5d169fabc751def39a7012
+size 3438167536

vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.32.0.dev0",
+  "_name_or_path": "/media/ubuntun/hdd/THj/model/stable-diffusion-v1-4",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fbcf0ebe55a0984f5a5e00d8c4521d52359af7229bb4d81890039d2aa16dd7c
+size 167335342