Add files using upload-large-folder tool

README.md

@@ -0,0 +1,127 @@
+---
+license: apache-2.0
+library_name: diffusers
+tags:
+  - hsigene
+  - hyperspectral
+  - latent-diffusion
+  - controlnet
+  - arxiv:2409.12470
+pipeline_tag: image-to-image
+---
+
+# BiliSakura/HSIGene
+
+**Hyperspectral image generation** — HSIGene converted to diffusers format. Conditional generation with local controls (HED, MLSD, sketch, segmentation), global controls (content, text), and metadata embeddings. Outputs 48-band hyperspectral images (256×256 pixels).
+
+> Source: [HSIGene](https://arxiv.org/abs/2409.12470). Converted to diffusers format; model dir is self-contained (no external project for inference).
+
+## Conversion
+
+The main diffusion checkpoint (`last.ckpt`) must be downloaded from [GoogleDrive](https://drive.google.com/file/d/1euJAbsxCgG1wIu_Eh5nPfmiSP9suWsR4/view?usp=drive_link) and placed in `projects/HSIGene-Diffusers/checkpoints/`.
+
+**Note:** `models/raw/HSIGene` contains annotator/auxiliary models (body pose, depth, SAM, etc.) only — not the main diffusion checkpoint.
+
+```bash
+cd projects/HSIGene-Diffusers
+python convert_to_diffusers.py \
+    --config_path configs/inference.yaml \
+    --ckpt_path checkpoints/last.ckpt \
+    --output_dir /root/worksapce/models/BiliSakura/HSIGene
+```
+
+## Repository Structure (after conversion)
+
+| Component               | Path                     |
+|------------------------|--------------------------|
+| UNet (LocalControlUNet)| `unet/`                  |
+| VAE                    | `vae/`                   |
+| Text encoder (CLIP)    | `text_encoder/`          |
+| Local adapter         | `local_adapter/`         |
+| Global content adapter| `global_content_adapter/`|
+| Global text adapter    | `global_text_adapter/`   |
+| Metadata encoder       | `metadata_encoder/`      |
+| Scheduler              | `scheduler/`             |
+| Pipeline               | `pipeline_hsigene.py`    |
+| Config                 | `model_index.json`       |
+
+## Usage
+
+**Option 1 – No `sys.path.insert` (AeroGen-style):** Load the pipeline from the model path via `importlib`; the model dir is added to the path automatically.
+
+```python
+import importlib.util
+import sys
+
+model_path = "/path/to/HSIGene"  # or "BiliSakura/HSIGene" for Hub
+spec = importlib.util.spec_from_file_location("pipeline_hsigene", f"{model_path}/pipeline_hsigene.py")
+mod = importlib.util.module_from_spec(spec)
+sys.modules["pipeline_hsigene"] = mod
+spec.loader.exec_module(mod)
+
+pipe = mod.HSIGenePipeline.from_pretrained(model_path)
+pipe = pipe.to("cuda")
+```
+
+**Option 2 – With `sys.path.insert`:** Simpler if you are fine adding the model dir to the path once.
+
+```python
+import sys
+sys.path.insert(0, "/path/to/HSIGene")
+from pipeline_hsigene import HSIGenePipeline
+
+pipe = HSIGenePipeline.from_pretrained("/path/to/HSIGene")
+pipe = pipe.to("cuda")
+```
+
+**Option 3 – `DiffusionPipeline.from_pretrained`:** May work with `trust_remote_code=True`. If you see "raw config (list)" errors (e.g. when loading from cache), use Option 1 or 2 instead.
+
+```python
+from diffusers import DiffusionPipeline
+pipe = DiffusionPipeline.from_pretrained("/path/to/HSIGene", trust_remote_code=True)
+pipe = pipe.to("cuda")
+```
+
+**Dependencies:** `pip install diffusers transformers torch einops safetensors`
+
+```python
+# Conditional generation
+output = pipe(
+    prompt="Wasteland",
+    num_samples=1,
+    height=256,
+    width=256,
+    num_inference_steps=50,
+    local_conditions=local_tensor,   # (B, 18, H, W) or None
+    global_conditions=global_tensor, # (B, 768) or None
+    metadata=metadata_tensor,        # (7,) or (B, 7) or None
+    guidance_scale=1.0,
+)
+images = output.images  # (B, H, W, 48) in [0, 1]
+```
+
+### Conditioning
+
+- **Local**: 18-channel maps (HED, MLSD, sketch, segmentation, etc.) at 512×512 default.
+- **Global**: 768-dim CLIP features from reference images.
+- **Metadata**: 7-dim vector.
+- **Text**: Via `prompt`; use `text_strength` to scale.
+
+## Model Sources
+
+- **Paper**: [HSIGene: A Foundation Model For Hyperspectral Image Generation](https://arxiv.org/abs/2409.12470)
+- **Checkpoint**: [GoogleDrive](https://drive.google.com/file/d/1euJAbsxCgG1wIu_Eh5nPfmiSP9suWsR4/view?usp=drive_link)
+- **Annotators**: [BaiduNetdisk](https://pan.baidu.com/s/1K1Y__blA6uJVV9l1QG7QvQ?pwd=98f1) (code: 98f1) → `data_prepare/annotator/ckpts`
+
+## Citation
+
+```bibtex
+@misc{pang2024hsigenefoundationmodelhyperspectral,
+  title={HSIGene: A Foundation Model For Hyperspectral Image Generation},
+  author={Li Pang and Datao Tang and Shuang Xu and Deyu Meng and Xiangyong Cao},
+  year={2024},
+  eprint={2409.12470},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV},
+}
+```

global_content_adapter/__init__.py

@@ -0,0 +1,5 @@
+"""Global content adapter for HSIGene."""
+
+from .model import GlobalContentAdapter
+
+__all__ = ["GlobalContentAdapter"]

global_content_adapter/config.json

@@ -0,0 +1,8 @@
+{
+  "_target": "hsigene.GlobalContentAdapter",
+  "in_dim": 768,
+  "channel_mult": [
+    2,
+    4
+  ]
+}

global_content_adapter/model.py

@@ -0,0 +1,56 @@
+"""GlobalContentAdapter - FFN-based adapter for global content conditioning."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
+        )
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class GlobalContentAdapter(nn.Module):
+    def __init__(self, in_dim, channel_mult=None):
+        super().__init__()
+        channel_mult = channel_mult or [2, 4]
+        dim_out1, mult1 = in_dim * channel_mult[0], channel_mult[0] * 2
+        dim_out2, mult2 = in_dim * channel_mult[1], channel_mult[1] * 2 // channel_mult[0]
+        self.in_dim = in_dim
+        self.channel_mult = channel_mult
+        self.ff1 = FeedForward(in_dim, dim_out=dim_out1, mult=mult1, glu=True, dropout=0.0)
+        self.ff2 = FeedForward(dim_out1, dim_out=dim_out2, mult=mult2, glu=True, dropout=0.0)
+        self.norm1 = nn.LayerNorm(in_dim)
+        self.norm2 = nn.LayerNorm(dim_out1)
+
+    def forward(self, x):
+        x = self.ff1(self.norm1(x))
+        x = self.ff2(self.norm2(x))
+        x = rearrange(x, "b (n d) -> b n d", n=self.channel_mult[-1], d=self.in_dim).contiguous()
+        return x

global_text_adapter/__init__.py

@@ -0,0 +1,5 @@
+"""Global text adapter for HSIGene."""
+
+from .model import GlobalTextAdapter
+
+__all__ = ["GlobalTextAdapter"]

global_text_adapter/config.json

@@ -0,0 +1,4 @@
+{
+  "_target": "hsigene.GlobalTextAdapter",
+  "in_dim": 768
+}

global_text_adapter/model.py

@@ -0,0 +1,52 @@
+"""GlobalTextAdapter - FFN-based adapter for global text conditioning."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
+        )
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class GlobalTextAdapter(nn.Module):
+    def __init__(self, in_dim, max_len=768):
+        super().__init__()
+        self.in_dim = in_dim
+        dim_out1 = in_dim * 2
+        dim_out2 = in_dim
+        self.ff1 = FeedForward(in_dim, dim_out=dim_out1, mult=2, glu=True, dropout=0.0)
+        self.ff2 = FeedForward(dim_out1, dim_out=dim_out2, mult=4, glu=True, dropout=0.0)
+        self.norm1 = nn.LayerNorm(in_dim)
+        self.norm2 = nn.LayerNorm(dim_out1)
+
+    def forward(self, x):
+        x = self.ff1(self.norm1(x))
+        x = self.ff2(self.norm2(x))
+        return x

local_adapter/attention.py

@@ -0,0 +1,271 @@
+"""HSIGene attention modules - FeedForward, CrossAttention, SpatialTransformer."""
+
+from inspect import isfunction
+from typing import Optional, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import einsum
+
+from .utils import checkpoint, zero_module, exists
+
+try:
+    import xformers
+    import xformers.ops
+    XFORMERS_IS_AVAILABLE = True
+except ImportError:
+    XFORMERS_IS_AVAILABLE = False
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+import os
+_ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
+
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
+        )
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+def Normalize(in_channels, num_groups=32):
+    return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x, context=None, mask=None):
+        h = self.heads
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
+        if _ATTN_PRECISION == "fp32":
+            with torch.autocast(enabled=False, device_type="cuda"):
+                q, k = q.float(), k.float()
+                sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
+        else:
+            sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
+        del q, k
+        if exists(mask):
+            mask = rearrange(mask, "b ... -> b (...)")
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, "b j -> (b h) () j", h=h)
+            sim.masked_fill_(~mask, max_neg_value)
+        sim = sim.softmax(dim=-1)
+        out = einsum("b i j, b j d -> b i d", sim, v)
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
+        return self.to_out(out)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+        self.heads = heads
+        self.dim_head = dim_head
+        self.scale = dim_head ** -0.5
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim),
+            nn.Dropout(dropout),
+        )
+        self.attention_op: Optional[Any] = None
+
+    def forward(self, x, context=None, mask=None):
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+        if XFORMERS_IS_AVAILABLE:
+            out = xformers.ops.memory_efficient_attention(
+                q, k, v, attn_bias=None, op=self.attention_op
+            )
+        else:
+            sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
+            sim = sim.softmax(dim=-1)
+            out = torch.einsum("b i j, b j d -> b i d", sim, v)
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,
+        "softmax-xformers": MemoryEfficientCrossAttention,
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        disable_self_attn=False,
+    ):
+        super().__init__()
+        attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILABLE else "softmax"
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim if self.disable_self_attn else None,
+        )
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = attn_cls(
+            query_dim=dim,
+            context_dim=context_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+        )
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None):
+        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None):
+        x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
+        x = self.attn2(self.norm2(x), context=context) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        n_heads,
+        d_head,
+        depth=1,
+        dropout=0.0,
+        context_dim=None,
+        disable_self_attn=False,
+        use_linear=False,
+        use_checkpoint=True,
+    ):
+        super().__init__()
+        if exists(context_dim) and not isinstance(context_dim, list):
+            context_dim = [context_dim]
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+        if not use_linear:
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    dropout=dropout,
+                    context_dim=context_dim[d] if isinstance(context_dim, list) else context_dim,
+                    disable_self_attn=disable_self_attn,
+                    checkpoint=use_checkpoint,
+                )
+                for d in range(depth)
+            ]
+        )
+        if not use_linear:
+            self.proj_out = zero_module(
+                nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+            )
+        else:
+            self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+        self.use_linear = use_linear
+
+    def forward(self, x, context=None):
+        if not isinstance(context, list):
+            context = [context]
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            ctx = context[i] if i < len(context) else context[0]
+            x = block(x, context=ctx)
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in

local_adapter/config.json

@@ -0,0 +1,36 @@
+{
+  "_target": "hsigene.LocalAdapter",
+  "in_channels": 4,
+  "model_channels": 320,
+  "local_channels": 18,
+  "inject_channels": [
+    192,
+    256,
+    384,
+    512
+  ],
+  "inject_layers": [
+    1,
+    4,
+    7,
+    10
+  ],
+  "num_res_blocks": 2,
+  "attention_resolutions": [
+    4,
+    2,
+    1
+  ],
+  "channel_mult": [
+    1,
+    2,
+    4,
+    4
+  ],
+  "use_checkpoint": true,
+  "num_heads": 8,
+  "use_spatial_transformer": true,
+  "transformer_depth": 1,
+  "context_dim": 768,
+  "legacy": false
+}

local_adapter/diffusion.py

@@ -0,0 +1,608 @@
+"""HSIGene diffusion modules - UNet, ResBlock, etc. From openaimodel."""
+
+from abc import abstractmethod
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import (
+    checkpoint,
+    conv_nd,
+    linear,
+    zero_module,
+    normalization,
+    timestep_embedding,
+    exists,
+)
+from .attention import SpatialTransformer
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """Create a 1D, 2D, or 3D average pooling module."""
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def convert_module_to_f16(x):
+    pass
+
+
+def convert_module_to_f32(x):
+    pass
+
+
+class TimestepBlock(nn.Module):
+    """Any module where forward() takes timestep embeddings as a second argument."""
+
+    @abstractmethod
+    def forward(self, x, emb):
+        """Apply the module to `x` given `emb` timestep embeddings."""
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    """Sequential module that passes timestep embeddings to children that support it."""
+
+    def forward(self, x, emb, context=None):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            elif isinstance(layer, SpatialTransformer):
+                x = layer(x, context)
+            else:
+                x = layer(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """Upsampling layer with optional convolution."""
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        if self.dims == 3:
+            x = F.interpolate(
+                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            )
+        else:
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """Downsampling layer with optional convolution."""
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        if use_conv:
+            self.op = conv_nd(
+                dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(TimestepBlock):
+    """Residual block with timestep conditioning."""
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        use_checkpoint=False,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 3, padding=1)
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        return checkpoint(self._forward, (x, emb), self.parameters(), self.use_checkpoint)
+
+    def _forward(self, x, emb):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = emb_out.chunk(2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AttentionBlock(nn.Module):
+    """Spatial self-attention block."""
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        use_checkpoint=False,
+        use_new_attention_order=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert channels % num_head_channels == 0
+            self.num_heads = channels // num_head_channels
+        self.use_checkpoint = use_checkpoint
+        self.norm = normalization(channels)
+        self.qkv = conv_nd(1, channels, channels * 3, 1)
+        self.attention = (
+            QKVAttention(self.num_heads)
+            if use_new_attention_order
+            else QKVAttentionLegacy(self.num_heads)
+        )
+        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
+
+    def forward(self, x):
+        return checkpoint(self._forward, (x,), self.parameters(), True)
+
+    def _forward(self, x):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+class QKVAttentionLegacy(nn.Module):
+    """QKV attention - split heads before split qkv."""
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum("bct,bcs->bts", q * scale, k * scale)
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = torch.einsum("bts,bcs->bct", weight, v)
+        return a.reshape(bs, -1, length)
+
+
+class QKVAttention(nn.Module):
+    """QKV attention - split qkv before split heads."""
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.chunk(3, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum(
+            "bct,bcs->bts",
+            (q * scale).view(bs * self.n_heads, ch, length),
+            (k * scale).view(bs * self.n_heads, ch, length),
+        )
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = torch.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+        return a.reshape(bs, -1, length)
+
+
+class UNetModel(nn.Module):
+    """Full UNet with attention and timestep embedding."""
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=-1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+        use_spatial_transformer=False,
+        transformer_depth=1,
+        context_dim=None,
+        n_embed=None,
+        legacy=True,
+        disable_self_attentions=None,
+        num_attention_blocks=None,
+        disable_middle_self_attn=False,
+        use_linear_in_transformer=False,
+    ):
+        super().__init__()
+        if use_spatial_transformer:
+            assert context_dim is not None
+        if context_dim is not None:
+            assert use_spatial_transformer
+            if hasattr(context_dim, "__iter__") and not isinstance(context_dim, (list, tuple)):
+                context_dim = list(context_dim)
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+        if num_heads == -1:
+            assert num_head_channels != -1
+        if num_head_channels == -1:
+            assert num_heads != -1
+
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        if isinstance(num_res_blocks, int):
+            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+        else:
+            assert len(num_res_blocks) == len(channel_mult)
+            self.num_res_blocks = num_res_blocks
+
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.use_checkpoint = use_checkpoint
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.predict_codebook_ids = n_embed is not None
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        if num_classes is not None:
+            if isinstance(num_classes, int):
+                self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+            elif num_classes == "continuous":
+                self.label_emb = nn.Linear(1, time_embed_dim)
+            else:
+                raise ValueError()
+
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))]
+        )
+        self._feature_size = model_channels
+        input_block_chans = [model_channels]
+        ch = model_channels
+        ds = 1
+
+        for level, mult in enumerate(channel_mult):
+            for nr in range(self.num_res_blocks[level]):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=mult * model_channels,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads_cur = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+                    disabled_sa = (
+                        disable_self_attentions[level]
+                        if exists(disable_self_attentions)
+                        else False
+                    )
+                    if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
+                        attn_block = (
+                            AttentionBlock(
+                                ch,
+                                use_checkpoint=use_checkpoint,
+                                num_heads=num_heads,
+                                num_head_channels=dim_head,
+                                use_new_attention_order=use_new_attention_order,
+                            )
+                            if not use_spatial_transformer
+                            else SpatialTransformer(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                depth=transformer_depth,
+                                context_dim=context_dim,
+                                disable_self_attn=disabled_sa,
+                                use_linear=use_linear_in_transformer,
+                                use_checkpoint=use_checkpoint,
+                            )
+                        )
+                        layers.append(attn_block)
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                down_block = (
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=out_ch,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                        down=True,
+                    )
+                    if resblock_updown
+                    else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                )
+                self.input_blocks.append(TimestepEmbedSequential(down_block))
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        if num_head_channels == -1:
+            dim_head = ch // num_heads
+        else:
+            num_heads_cur = ch // num_head_channels
+            dim_head = num_head_channels
+        if legacy:
+            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+        mid_attn = (
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=dim_head,
+                use_new_attention_order=use_new_attention_order,
+            )
+            if not use_spatial_transformer
+            else SpatialTransformer(
+                ch,
+                num_heads,
+                dim_head,
+                depth=transformer_depth,
+                context_dim=context_dim,
+                disable_self_attn=disable_middle_self_attn,
+                use_linear=use_linear_in_transformer,
+                use_checkpoint=use_checkpoint,
+            )
+        )
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            mid_attn,
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(self.num_res_blocks[level] + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=model_channels * mult,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = model_channels * mult
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads_cur = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        dim_head = (
+                            ch // num_heads if use_spatial_transformer else num_head_channels
+                        )
+                    disabled_sa = (
+                        disable_self_attentions[level]
+                        if exists(disable_self_attentions)
+                        else False
+                    )
+                    if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
+                        attn_block = (
+                            AttentionBlock(
+                                ch,
+                                use_checkpoint=use_checkpoint,
+                                num_heads=num_heads_upsample,
+                                num_head_channels=dim_head,
+                                use_new_attention_order=use_new_attention_order,
+                            )
+                            if not use_spatial_transformer
+                            else SpatialTransformer(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                depth=transformer_depth,
+                                context_dim=context_dim,
+                                disable_self_attn=disabled_sa,
+                                use_linear=use_linear_in_transformer,
+                                use_checkpoint=use_checkpoint,
+                            )
+                        )
+                        layers.append(attn_block)
+                if level and i == self.num_res_blocks[level]:
+                    out_ch = ch
+                    up_block = (
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    layers.append(up_block)
+                    ds //= 2
+                self.output_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            normalization(ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+        )
+        if self.predict_codebook_ids:
+            self.id_predictor = nn.Sequential(
+                normalization(ch),
+                conv_nd(dims, model_channels, n_embed, 1),
+            )
+
+    def forward(self, x, timesteps=None, metadata=None, context=None, y=None, **kwargs):
+        assert (y is not None) == (self.num_classes is not None)
+        hs = []
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        emb = self.time_embed(t_emb)
+        if metadata is not None:
+            if isinstance(metadata, (list, tuple)) and len(metadata) == 1:
+                metadata = metadata[0]
+            emb = emb + metadata
+
+        if self.num_classes is not None:
+            assert y.shape[0] == x.shape[0]
+            emb = emb + self.label_emb(y)
+
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb, context)
+            hs.append(h)
+        h = self.middle_block(h, emb, context)
+        for module in self.output_blocks:
+            h = torch.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context)
+        h = h.type(x.dtype)
+        if self.predict_codebook_ids:
+            return self.id_predictor(h)
+        return self.out(h)

local_adapter/model.py

@@ -0,0 +1,435 @@
+"""HSIGene adapters - LocalAdapter, LocalControlUNetModel, GlobalContentAdapter, etc."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from .utils import (
+    checkpoint,
+    conv_nd,
+    linear,
+    zero_module,
+    timestep_embedding,
+    exists,
+)
+from .attention import SpatialTransformer
+from .diffusion import (
+    TimestepBlock,
+    TimestepEmbedSequential,
+    ResBlock,
+    Downsample,
+    AttentionBlock,
+)
+
+
+class LocalTimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    """Sequential that handles LocalResBlock, TimestepBlock, SpatialTransformer."""
+
+    def forward(self, x, emb, context=None, local_features=None):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            elif isinstance(layer, SpatialTransformer):
+                x = layer(x, context)
+            elif isinstance(layer, LocalResBlock):
+                x = layer(x, emb, local_features)
+            else:
+                x = layer(x)
+        return x
+
+
+class FDN(nn.Module):
+    def __init__(self, norm_nc, label_nc):
+        super().__init__()
+        ks = 3
+        pw = ks // 2
+        self.param_free_norm = nn.GroupNorm(32, norm_nc, affine=False)
+        self.conv_gamma = nn.Conv2d(label_nc, norm_nc, kernel_size=ks, padding=pw)
+        self.conv_beta = nn.Conv2d(label_nc, norm_nc, kernel_size=ks, padding=pw)
+
+    def forward(self, x, local_features):
+        normalized = self.param_free_norm(x)
+        assert local_features.size()[2:] == x.size()[2:]
+        gamma = self.conv_gamma(local_features)
+        beta = self.conv_beta(local_features)
+        return normalized * (1 + gamma) + beta
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, in_dim):
+        super().__init__()
+        self.query_conv = nn.Conv2d(in_dim, in_dim // 8, kernel_size=1)
+        self.key_conv = nn.Conv2d(in_dim, in_dim // 8, kernel_size=1)
+        self.value_conv = nn.Conv2d(in_dim, in_dim, kernel_size=1)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x):
+        batch, C, width, height = x.size()
+        query = self.query_conv(x).view(batch, -1, width * height).permute(0, 2, 1)
+        key = self.key_conv(x).view(batch, -1, width * height)
+        value = self.value_conv(x).view(batch, -1, width * height)
+        attention = self.softmax(torch.bmm(query, key))
+        out = torch.bmm(value, attention.permute(0, 2, 1))
+        out = out.view(batch, C, width, height)
+        return out + x
+
+
+class EnhancedFDN(nn.Module):
+    def __init__(self, norm_nc, label_nc):
+        super().__init__()
+        self.fdn = FDN(norm_nc, label_nc)
+        self.attention = SelfAttention(norm_nc)
+
+    def forward(self, x, local_features):
+        x = self.attention(x)
+        return self.fdn(x, local_features)
+
+
+class LocalResBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        dims=2,
+        use_checkpoint=False,
+        inject_channels=None,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        self.norm_in = EnhancedFDN(channels, inject_channels)
+        self.norm_out = EnhancedFDN(self.out_channels, inject_channels)
+
+        self.in_layers = nn.Sequential(
+            nn.Identity(),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            linear(emb_channels, self.out_channels),
+        )
+        self.out_layers = nn.Sequential(
+            nn.Identity(),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb, local_conditions):
+        return checkpoint(
+            self._forward, (x, emb, local_conditions), self.parameters(), self.use_checkpoint
+        )
+
+    def _forward(self, x, emb, local_conditions):
+        h = self.norm_in(x, local_conditions)
+        h = self.in_layers(h)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        h = h + emb_out
+        h = self.norm_out(h, local_conditions)
+        h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class FeatureExtractor(nn.Module):
+    def __init__(self, local_channels, inject_channels, dims=2):
+        super().__init__()
+        self.pre_extractor = LocalTimestepEmbedSequential(
+            conv_nd(dims, local_channels, 32, 3, padding=1),
+            nn.SiLU(),
+            conv_nd(dims, 32, 64, 3, padding=1, stride=2),
+            nn.SiLU(),
+            conv_nd(dims, 64, 64, 3, padding=1),
+            nn.SiLU(),
+            conv_nd(dims, 64, 128, 3, padding=1, stride=2),
+            nn.SiLU(),
+            conv_nd(dims, 128, 128, 3, padding=1),
+            nn.SiLU(),
+        )
+        self.extractors = nn.ModuleList([
+            LocalTimestepEmbedSequential(
+                conv_nd(dims, 128, inject_channels[0], 3, padding=1, stride=2),
+                nn.SiLU(),
+            ),
+            LocalTimestepEmbedSequential(
+                conv_nd(dims, inject_channels[0], inject_channels[1], 3, padding=1, stride=2),
+                nn.SiLU(),
+            ),
+            LocalTimestepEmbedSequential(
+                conv_nd(dims, inject_channels[1], inject_channels[2], 3, padding=1, stride=2),
+                nn.SiLU(),
+            ),
+            LocalTimestepEmbedSequential(
+                conv_nd(dims, inject_channels[2], inject_channels[3], 3, padding=1, stride=2),
+                nn.SiLU(),
+            ),
+        ])
+        self.zero_convs = nn.ModuleList([
+            zero_module(conv_nd(dims, inject_channels[0], inject_channels[0], 3, padding=1)),
+            zero_module(conv_nd(dims, inject_channels[1], inject_channels[1], 3, padding=1)),
+            zero_module(conv_nd(dims, inject_channels[2], inject_channels[2], 3, padding=1)),
+            zero_module(conv_nd(dims, inject_channels[3], inject_channels[3], 3, padding=1)),
+        ])
+
+    def forward(self, local_conditions):
+        local_features = self.pre_extractor(local_conditions, None)
+        output_features = []
+        for idx in range(len(self.extractors)):
+            local_features = self.extractors[idx](local_features, None)
+            output_features.append(self.zero_convs[idx](local_features))
+        return output_features
+
+
+class LocalAdapter(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        local_channels,
+        inject_channels,
+        inject_layers,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=-1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+        use_spatial_transformer=False,
+        transformer_depth=1,
+        context_dim=None,
+        n_embed=None,
+        legacy=True,
+        disable_self_attentions=None,
+        num_attention_blocks=None,
+        disable_middle_self_attn=False,
+        use_linear_in_transformer=False,
+    ):
+        super().__init__()
+        if context_dim is not None:
+            if hasattr(context_dim, "__iter__") and not isinstance(context_dim, (list, tuple)):
+                context_dim = list(context_dim)
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        self.dims = dims
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.inject_layers = inject_layers
+        if isinstance(num_res_blocks, int):
+            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+        else:
+            assert len(num_res_blocks) == len(channel_mult)
+            self.num_res_blocks = num_res_blocks
+
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.use_checkpoint = use_checkpoint
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.predict_codebook_ids = n_embed is not None
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        self.feature_extractor = FeatureExtractor(local_channels, inject_channels)
+        self.input_blocks = nn.ModuleList([
+            LocalTimestepEmbedSequential(
+                conv_nd(dims, in_channels, model_channels, 3, padding=1)
+            ),
+        ])
+        self.zero_convs = nn.ModuleList([self._make_zero_conv(model_channels)])
+
+        self._feature_size = model_channels
+        input_block_chans = [model_channels]
+        ch = model_channels
+        ds = 1
+
+        for level, mult in enumerate(channel_mult):
+            for nr in range(self.num_res_blocks[level]):
+                if (1 + 3 * level + nr) in self.inject_layers:
+                    layers = [
+                        LocalResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=mult * model_channels,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            inject_channels=inject_channels[level],
+                        )
+                    ]
+                else:
+                    layers = [
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=mult * model_channels,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                        )
+                    ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        dim_head = num_head_channels
+                    if legacy:
+                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+                    disabled_sa = (
+                        disable_self_attentions[level]
+                        if exists(disable_self_attentions)
+                        else False
+                    )
+                    if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
+                        block = (
+                            AttentionBlock(
+                                ch,
+                                use_checkpoint=use_checkpoint,
+                                num_heads=num_heads,
+                                num_head_channels=dim_head,
+                                use_new_attention_order=use_new_attention_order,
+                            )
+                            if not use_spatial_transformer
+                            else SpatialTransformer(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                depth=transformer_depth,
+                                context_dim=context_dim,
+                                disable_self_attn=disabled_sa,
+                                use_linear=use_linear_in_transformer,
+                                use_checkpoint=use_checkpoint,
+                            )
+                        )
+                        layers.append(block)
+                self.input_blocks.append(LocalTimestepEmbedSequential(*layers))
+                self.zero_convs.append(self._make_zero_conv(ch))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                down_block = (
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=out_ch,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                        down=True,
+                    )
+                    if resblock_updown
+                    else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                )
+                self.input_blocks.append(LocalTimestepEmbedSequential(down_block))
+                ch = out_ch
+                input_block_chans.append(ch)
+                self.zero_convs.append(self._make_zero_conv(ch))
+                ds *= 2
+                self._feature_size += ch
+
+        if num_head_channels == -1:
+            dim_head = ch // num_heads
+        else:
+            dim_head = num_head_channels
+        if legacy:
+            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+        mid_attn = (
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=dim_head,
+                use_new_attention_order=use_new_attention_order,
+            )
+            if not use_spatial_transformer
+            else SpatialTransformer(
+                ch,
+                num_heads,
+                dim_head,
+                depth=transformer_depth,
+                context_dim=context_dim,
+                disable_self_attn=disable_middle_self_attn,
+                use_linear=use_linear_in_transformer,
+                use_checkpoint=use_checkpoint,
+            )
+        )
+        self.middle_block = LocalTimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            mid_attn,
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self.middle_block_out = self._make_zero_conv(ch)
+        self._feature_size += ch
+
+    def _make_zero_conv(self, channels):
+        return LocalTimestepEmbedSequential(
+            zero_module(conv_nd(self.dims, channels, channels, 1, padding=0))
+        )
+
+    def forward(self, x, timesteps, context, local_conditions, **kwargs):
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        emb = self.time_embed(t_emb)
+        local_features = self.feature_extractor(local_conditions)
+
+        outs = []
+        h = x.type(self.dtype)
+        for layer_idx, (module, zero_conv) in enumerate(zip(self.input_blocks, self.zero_convs)):
+            if layer_idx in self.inject_layers:
+                feat_idx = self.inject_layers.index(layer_idx)
+                h = module(h, emb, context, local_features[feat_idx])
+            else:
+                h = module(h, emb, context)
+            outs.append(zero_conv(h, emb, context))
+
+        h = self.middle_block(h, emb, context)
+        outs.append(self.middle_block_out(h, emb, context))
+        return outs

local_adapter/utils.py

@@ -0,0 +1,90 @@
+"""HSIGene utilities - no ldm/models imports."""
+
+import math
+import torch
+import torch.nn as nn
+from einops import repeat
+
+
+def exists(val):
+    return val is not None
+
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    if repeat_only:
+        return repeat(timesteps, "b -> b d", d=dim)
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+def conv_nd(dims, *args, **kwargs):
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+    return nn.Linear(*args, **kwargs)
+
+
+def zero_module(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def checkpoint(func, inputs, params, flag):
+    if flag:
+        return _CheckpointFunction.apply(func, len(inputs), *(tuple(inputs) + tuple(params)))
+    return func(*inputs)
+
+
+class _CheckpointFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+        ctx.gpu_autocast_kwargs = {
+            "enabled": torch.is_autocast_enabled(),
+            "dtype": torch.get_autocast_gpu_dtype(),
+            "cache_enabled": torch.is_autocast_cache_enabled(),
+        }
+        with torch.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+        with torch.enable_grad(), torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = torch.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        return (None, None) + input_grads
+
+
+def normalization(channels):
+    return GroupNorm32(32, channels)
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+

metadata_encoder/__init__.py

@@ -0,0 +1,5 @@
+"""Metadata encoder for HSIGene."""
+
+from .model import MetadataEmbeddings, metadata_embeddings
+
+__all__ = ["MetadataEmbeddings", "metadata_embeddings"]

metadata_encoder/config.json

@@ -0,0 +1,7 @@
+{
+  "_target": "hsigene.metadata_embeddings",
+  "max_value": 1000,
+  "embedding_dim": 320,
+  "metadata_dim": 7,
+  "max_period": 10000
+}

metadata_encoder/model.py

@@ -0,0 +1,77 @@
+"""Metadata embeddings - SinusoidalEmbedding + MLPs for metadata conditioning."""
+
+import torch
+import torch.nn as nn
+
+
+class SinusoidalEmbedding(nn.Module):
+    """Sinusoidal embedding for metadata."""
+
+    def __init__(self, max_value, embedding_dim):
+        super().__init__()
+        self.max_value = max_value
+        self.embedding_dim = embedding_dim
+        self.omega = 10000.0
+
+    def forward(self, k):
+        device = k.device
+        k_normalized = k * self.max_value
+        embedding = torch.zeros(
+            (k.size(0), k.size(1), self.embedding_dim),
+            device=device,
+            dtype=k.dtype,
+        )
+        for j in range(k.size(1)):
+            for i in range(self.embedding_dim // 2):
+                omega_term = self.omega ** (-2 * i / self.embedding_dim)
+                embedding[:, j, 2 * i] = torch.sin(k_normalized[:, j] * omega_term)
+                embedding[:, j, 2 * i + 1] = torch.cos(k_normalized[:, j] * omega_term)
+        return embedding.view(k.size(0), -1)
+
+
+def create_condition_vector(embedded_metadata, mlp_models, embedding_dim):
+    """Create condition vector from metadata embeddings and MLPs."""
+    metadata_embeddings = [
+        mlp_models[j](embedded_metadata[:, j * embedding_dim : (j + 1) * embedding_dim])
+        for j in range(len(mlp_models))
+    ]
+    return sum(metadata_embeddings)
+
+
+class MetadataMLP(nn.Module):
+    def __init__(self, input_dim, embedding_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, embedding_dim)
+
+    def forward(self, x):
+        return self.fc1(x)
+
+
+class MetadataEmbeddings(nn.Module):
+    """Metadata embeddings - SinusoidalEmbedding + MLPs."""
+
+    def __init__(self, max_value, embedding_dim, max_period, metadata_dim):
+        super().__init__()
+        self.sinusoidal_embedding = SinusoidalEmbedding(max_value, embedding_dim)
+        self.mlp_models = nn.ModuleList([
+            MetadataMLP(embedding_dim, embedding_dim * 4)
+            for _ in range(metadata_dim)
+        ])
+        self.max_period = max_period
+        self.embedding_dim = embedding_dim
+        self.metadata_dim = metadata_dim
+        self.max_value = max_value
+
+    def forward(self, metadata=None):
+        while isinstance(metadata, (list, tuple)) and len(metadata) == 1:
+            metadata = metadata[0]
+        if metadata.dim() == 1:
+            metadata = metadata.unsqueeze(0)
+        embedded_metadata = self.sinusoidal_embedding(metadata)
+        return create_condition_vector(
+            embedded_metadata, self.mlp_models, self.embedding_dim
+        )
+
+
+# Alias for config compatibility
+metadata_embeddings = MetadataEmbeddings

model_index.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": ["pipeline_hsigene", "HSIGenePipeline"],
+  "_diffusers_version": "0.25.0",
+  "scheduler": ["diffusers", "DDIMScheduler"],
+  "unet": ["pipeline_hsigene", "HSIGenePipeline"],
+  "vae": ["pipeline_hsigene", "HSIGenePipeline"],
+  "text_encoder": ["pipeline_hsigene", "HSIGenePipeline"],
+  "local_adapter": ["pipeline_hsigene", "HSIGenePipeline"],
+  "global_content_adapter": ["pipeline_hsigene", "HSIGenePipeline"],
+  "global_text_adapter": ["pipeline_hsigene", "HSIGenePipeline"],
+  "metadata_encoder": ["pipeline_hsigene", "HSIGenePipeline"],
+  "scale_factor": 0.18215,
+  "conditioning_key": "crossattn"
+}

modular_pipeline.py

@@ -0,0 +1,111 @@
+"""
+HSIGene modular components: path setup and component loading.
+
+AeroGen-style: ensure_ldm_path adds model dir to sys.path so hsigene can be imported.
+No manual sys.path.insert needed when using DiffusionPipeline.from_pretrained(path).
+"""
+
+import importlib
+import json
+import sys
+from pathlib import Path
+from typing import Union
+
+from diffusers import DDIMScheduler
+
+# Ensure model dir is on path for hsigene imports
+_pipeline_dir = Path(__file__).resolve().parent
+if str(_pipeline_dir) not in sys.path:
+    sys.path.insert(0, str(_pipeline_dir))
+
+_COMPONENT_NAMES = (
+    "unet", "vae", "text_encoder", "local_adapter",
+    "global_content_adapter", "global_text_adapter", "metadata_encoder",
+)
+
+_TARGET_MAP = {
+    "hsigene_models.HSIGeneUNet": "unet.model.HSIGeneUNet",
+    "hsigene.HSIGeneUNet": "unet.model.HSIGeneUNet",
+    "hsigene_models.HSIGeneAutoencoderKL": "vae.model.HSIGeneAutoencoderKL",
+    "hsigene.HSIGeneAutoencoderKL": "vae.model.HSIGeneAutoencoderKL",
+    "ldm.modules.encoders.modules.FrozenCLIPEmbedder": "text_encoder.model.CLIPTextEncoder",
+    "hsigene.CLIPTextEncoder": "text_encoder.model.CLIPTextEncoder",
+    "models.local_adapter.LocalAdapter": "local_adapter.model.LocalAdapter",
+    "hsigene.LocalAdapter": "local_adapter.model.LocalAdapter",
+    "models.global_adapter.GlobalContentAdapter": "global_content_adapter.model.GlobalContentAdapter",
+    "hsigene.GlobalContentAdapter": "global_content_adapter.model.GlobalContentAdapter",
+    "models.global_adapter.GlobalTextAdapter": "global_text_adapter.model.GlobalTextAdapter",
+    "hsigene.GlobalTextAdapter": "global_text_adapter.model.GlobalTextAdapter",
+    "models.metadata_embedding.metadata_embeddings": "metadata_encoder.model.metadata_embeddings",
+    "hsigene.metadata_embeddings": "metadata_encoder.model.metadata_embeddings",
+}
+
+
+def ensure_ldm_path(pretrained_model_name_or_path: Union[str, Path]) -> Path:
+    """Add model repo to path so hsigene can be imported. Returns resolved path."""
+    path = Path(pretrained_model_name_or_path)
+    if not path.exists():
+        from huggingface_hub import snapshot_download
+        path = Path(snapshot_download(pretrained_model_name_or_path))
+    path = path.resolve()
+    s = str(path)
+    if s not in sys.path:
+        sys.path.insert(0, s)
+    return path
+
+
+def _get_class(target: str):
+    module_path, cls_name = target.rsplit(".", 1)
+    mod = importlib.import_module(module_path)
+    return getattr(mod, cls_name)
+
+
+def load_component(model_path: Path, name: str):
+    """Load a single component (unet, vae, text_encoder, etc.)."""
+    import torch
+    path = Path(model_path)
+    root = path.parent if path.name in _COMPONENT_NAMES and (path / "config.json").exists() else path
+    ensure_ldm_path(root)
+    comp_path = path if (path / "config.json").exists() and path.name in _COMPONENT_NAMES else path / name
+    with open(comp_path / "config.json") as f:
+        cfg = json.load(f)
+    target = cfg.pop("_target", None)
+    if not target:
+        raise ValueError(f"No _target in {comp_path / 'config.json'}")
+    target = _TARGET_MAP.get(target, target)
+    cls_ref = _get_class(target)
+    params = {k: v for k, v in cfg.items() if not k.startswith("_")}
+    comp = cls_ref(**params)
+    for wfile in ("diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.bin"):
+        wp = comp_path / wfile
+        if wp.exists():
+            if wfile.endswith(".safetensors"):
+                from safetensors.torch import load_file
+                state = load_file(str(wp))
+            else:
+                try:
+                    state = torch.load(wp, map_location="cpu", weights_only=True)
+                except TypeError:
+                    state = torch.load(wp, map_location="cpu")
+            comp.load_state_dict(state, strict=True)
+            break
+    comp.eval()
+    return comp
+
+
+def load_components(model_path: Union[str, Path]) -> dict:
+    """Load all pipeline components. Returns dict with components, scheduler, scale_factor."""
+    path = Path(ensure_ldm_path(model_path))
+    if path.name in _COMPONENT_NAMES and (path / "config.json").exists():
+        path = path.parent
+    scheduler = DDIMScheduler.from_pretrained(path / "scheduler")
+    components = {}
+    for name in _COMPONENT_NAMES:
+        components[name] = load_component(path, name)
+    scale_factor = 0.18215
+    if (path / "model_index.json").exists():
+        with open(path / "model_index.json") as f:
+            scale_factor = json.load(f).get("scale_factor", scale_factor)
+    components["scheduler"] = scheduler
+    components["scale_factor"] = scale_factor
+    return components

pipeline_hsigene.py

@@ -0,0 +1,468 @@
+"""HSIGenePipeline - diffusers DiffusionPipeline for HSIGene hyperspectral generation.
+
+AeroGen-style loading: use DiffusionPipeline.from_pretrained(path) - no sys.path.insert needed.
+Self-contained: loading logic inlined (no separate modular_pipeline import).
+"""
+
+import importlib
+import json
+import sys
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from dataclasses import dataclass
+
+from diffusers import DDIMScheduler, DiffusionPipeline
+from diffusers.utils import BaseOutput
+
+# Re-export for diffusers component loading (load_method lookup)
+DiffusionPipeline = DiffusionPipeline
+
+# Inline path/loading (AeroGen-style) - self-contained for diffusers cache loading
+_pipeline_dir = Path(__file__).resolve().parent
+if str(_pipeline_dir) not in sys.path:
+    sys.path.insert(0, str(_pipeline_dir))
+
+# Register as "pipeline_hsigene" so diffusers' get_class_obj_and_candidates finds us when it does
+# importlib.import_module("pipeline_hsigene") during component loading. (We may be loaded as
+# "diffusers_modules.local.xxx.pipeline_hsigene" from cache, so this alias is required.)
+sys.modules["pipeline_hsigene"] = sys.modules[__name__]
+
+_COMPONENT_NAMES = (
+    "unet", "vae", "text_encoder", "local_adapter",
+    "global_content_adapter", "global_text_adapter", "metadata_encoder",
+)
+
+_TARGET_MAP = {
+    "hsigene_models.HSIGeneUNet": "unet.model.HSIGeneUNet",
+    "hsigene.HSIGeneUNet": "unet.model.HSIGeneUNet",
+    "hsigene_models.HSIGeneAutoencoderKL": "vae.model.HSIGeneAutoencoderKL",
+    "hsigene.HSIGeneAutoencoderKL": "vae.model.HSIGeneAutoencoderKL",
+    "ldm.modules.encoders.modules.FrozenCLIPEmbedder": "text_encoder.model.CLIPTextEncoder",
+    "hsigene.CLIPTextEncoder": "text_encoder.model.CLIPTextEncoder",
+    "models.local_adapter.LocalAdapter": "local_adapter.model.LocalAdapter",
+    "hsigene.LocalAdapter": "local_adapter.model.LocalAdapter",
+    "models.global_adapter.GlobalContentAdapter": "global_content_adapter.model.GlobalContentAdapter",
+    "hsigene.GlobalContentAdapter": "global_content_adapter.model.GlobalContentAdapter",
+    "models.global_adapter.GlobalTextAdapter": "global_text_adapter.model.GlobalTextAdapter",
+    "hsigene.GlobalTextAdapter": "global_text_adapter.model.GlobalTextAdapter",
+    "models.metadata_embedding.metadata_embeddings": "metadata_encoder.model.metadata_embeddings",
+    "hsigene.metadata_embeddings": "metadata_encoder.model.metadata_embeddings",
+}
+
+
+def ensure_ldm_path(pretrained_model_name_or_path: Union[str, Path]) -> Path:
+    """Add model repo to path so hsigene can be imported. Returns resolved path."""
+    path = Path(pretrained_model_name_or_path)
+    if not path.exists():
+        from huggingface_hub import snapshot_download
+        path = Path(snapshot_download(pretrained_model_name_or_path))
+    path = path.resolve()
+    s = str(path)
+    if s not in sys.path:
+        sys.path.insert(0, s)
+    return path
+
+
+def _get_class(target: str):
+    module_path, cls_name = target.rsplit(".", 1)
+    mod = importlib.import_module(module_path)
+    return getattr(mod, cls_name)
+
+
+def load_component(model_path: Path, name: str):
+    """Load a single component (unet, vae, text_encoder, etc.)."""
+    path = Path(model_path)
+    root = path.parent if path.name in _COMPONENT_NAMES and (path / "config.json").exists() else path
+    ensure_ldm_path(root)
+    comp_path = path if (path / "config.json").exists() and path.name in _COMPONENT_NAMES else path / name
+    with open(comp_path / "config.json") as f:
+        cfg = json.load(f)
+    target = cfg.pop("_target", None)
+    if not target:
+        raise ValueError(f"No _target in {comp_path / 'config.json'}")
+    target = _TARGET_MAP.get(target, target)
+    cls_ref = _get_class(target)
+    params = {k: v for k, v in cfg.items() if not k.startswith("_")}
+    comp = cls_ref(**params)
+    for wfile in ("diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.bin"):
+        wp = comp_path / wfile
+        if wp.exists():
+            if wfile.endswith(".safetensors"):
+                from safetensors.torch import load_file
+                state = load_file(str(wp))
+            else:
+                try:
+                    state = torch.load(wp, map_location="cpu", weights_only=True)
+                except TypeError:
+                    state = torch.load(wp, map_location="cpu")
+            comp.load_state_dict(state, strict=True)
+            break
+    comp.eval()
+    return comp
+
+
+def load_components(model_path: Union[str, Path]) -> dict:
+    """Load all pipeline components."""
+    path = Path(ensure_ldm_path(model_path))
+    if path.name in _COMPONENT_NAMES and (path / "config.json").exists():
+        path = path.parent
+    scheduler = DDIMScheduler.from_pretrained(path / "scheduler")
+    components = {}
+    for name in _COMPONENT_NAMES:
+        components[name] = load_component(path, name)
+    scale_factor = 0.18215
+    if (path / "model_index.json").exists():
+        with open(path / "model_index.json") as f:
+            scale_factor = json.load(f).get("scale_factor", scale_factor)
+    components["scheduler"] = scheduler
+    components["scale_factor"] = scale_factor
+    return components
+
+
+class _CRSModelWrapper(torch.nn.Module):
+    """Wrapper that mimics CRSControlNet interface."""
+
+    def __init__(
+        self,
+        unet,
+        vae,
+        text_encoder,
+        local_adapter,
+        global_content_adapter,
+        global_text_adapter,
+        metadata_emb,
+        scale_factor=0.18215,
+        local_control_scales=None,
+    ):
+        super().__init__()
+        self.model = type("Model", (), {"diffusion_model": unet})()
+        self.first_stage_model = vae
+        self.cond_stage_model = text_encoder
+        self.local_adapter = local_adapter
+        self.global_content_adapter = global_content_adapter
+        self.global_text_adapter = global_text_adapter
+        self.metadata_emb = metadata_emb
+        self.scale_factor = scale_factor
+        self.local_control_scales = local_control_scales or [1.0] * 13
+
+    @torch.no_grad()
+    def get_learned_conditioning(self, prompts):
+        return self.cond_stage_model(prompts)
+
+    def apply_model(self, x_noisy, t, cond, metadata=None, global_strength=1.0, text_strength=1.0, **kwargs):
+        if metadata is None:
+            metadata = cond["metadata"]
+        metadata_emb = self.metadata_emb(metadata)
+        content_t = cond["global_control"][0]
+        global_control = self.global_content_adapter(content_t)
+        cond_txt = torch.cat(cond["c_crossattn"], 1)
+        cond_txt = self.global_text_adapter(cond_txt)
+        cond_txt = F.normalize(cond_txt, p=2, dim=-1) * text_strength
+        global_control = F.normalize(global_control, p=2, dim=-1) * global_strength
+        cond_txt = torch.cat([cond_txt, global_control], dim=1)
+        local_control = torch.cat(cond["local_control"], 1)
+        local_control = self.local_adapter(
+            x=x_noisy, timesteps=t, context=cond_txt, local_conditions=local_control
+        )
+        local_control = [c * s for c, s in zip(local_control, self.local_control_scales)]
+        return self.model.diffusion_model(
+            x=x_noisy,
+            timesteps=t,
+            metadata=metadata_emb,
+            context=cond_txt,
+            local_control=local_control,
+            meta=True,
+        )
+
+    def decode_first_stage(self, z):
+        z = (1.0 / self.scale_factor) * z
+        return self.first_stage_model.decode(z)
+
+    def low_vram_shift(self, is_diffusing):
+        if is_diffusing:
+            self.model.diffusion_model = self.model.diffusion_model.cuda()
+            self.local_adapter = self.local_adapter.cuda()
+            self.global_text_adapter = self.global_text_adapter.cuda()
+            self.global_content_adapter = self.global_content_adapter.cuda()
+            self.first_stage_model = self.first_stage_model.cpu()
+            self.cond_stage_model = self.cond_stage_model.cpu()
+        else:
+            self.model.diffusion_model = self.model.diffusion_model.cpu()
+            self.local_adapter = self.local_adapter.cpu()
+            self.global_text_adapter = self.global_text_adapter.cpu()
+            self.global_content_adapter = self.global_content_adapter.cpu()
+            self.first_stage_model = self.first_stage_model.cuda()
+            self.cond_stage_model = self.cond_stage_model.cuda()
+
+
+@dataclass
+class HSIGeneOutput(BaseOutput):
+    """Output class for HSIGene pipeline."""
+
+    images: Optional[np.ndarray] = None
+    latents: Optional[torch.Tensor] = None
+
+
+def _is_component_list(v):
+    """Check if value is raw config format [library, class_name]."""
+    return isinstance(v, (list, tuple)) and len(v) == 2 and isinstance(v[0], str) and isinstance(v[1], str)
+
+
+class HSIGenePipeline(DiffusionPipeline):
+    """Pipeline for HSIGene hyperspectral image generation.
+
+    AeroGen-style: load with DiffusionPipeline.from_pretrained(path) - no sys.path.insert.
+    """
+
+    def register_modules(self, **kwargs):
+        """Override to handle list-format component specs from diffusers config."""
+        for name, module in kwargs.items():
+            if module is None or (isinstance(module, (tuple, list)) and len(module) > 0 and module[0] is None):
+                self.register_to_config(**{name: (None, None)})
+                setattr(self, name, module)
+            elif _is_component_list(module):
+                self.register_to_config(**{name: (module[0], module[1])})
+                setattr(self, name, module)
+            else:
+                from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple
+                library, class_name = _fetch_class_library_tuple(module)
+                self.register_to_config(**{name: (library, class_name)})
+                setattr(self, name, module)
+
+    def __init__(
+        self,
+        unet=None,
+        vae=None,
+        text_encoder=None,
+        local_adapter=None,
+        global_content_adapter=None,
+        global_text_adapter=None,
+        metadata_encoder=None,
+        scheduler=None,
+        crs_model=None,
+        scale_factor=0.18215,
+    ):
+        super().__init__()
+        if crs_model is not None:
+            self.register_modules(crs_model=crs_model, scheduler=scheduler)
+        else:
+            if any(_is_component_list(x) for x in (unet, vae, text_encoder, local_adapter,
+                    global_content_adapter, global_text_adapter, metadata_encoder) if x is not None):
+                raise ValueError(
+                    "HSIGene received raw config (list) instead of loaded components. "
+                    "Use HSIGenePipeline.from_pretrained(path) directly, or ensure the model "
+                    "directory (with hsigene package) is on the path when loading."
+                )
+            crs_model = _CRSModelWrapper(
+                unet=unet,
+                vae=vae,
+                text_encoder=text_encoder,
+                local_adapter=local_adapter,
+                global_content_adapter=global_content_adapter,
+                global_text_adapter=global_text_adapter,
+                metadata_emb=metadata_encoder,
+                scale_factor=scale_factor,
+            )
+            self.register_modules(
+                unet=unet,
+                vae=vae,
+                text_encoder=text_encoder,
+                local_adapter=local_adapter,
+                global_content_adapter=global_content_adapter,
+                global_text_adapter=global_text_adapter,
+                metadata_encoder=metadata_encoder,
+                scheduler=scheduler,
+                crs_model=crs_model,
+            )
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+        device: Optional[Union[str, torch.device]] = None,
+        subfolder: Optional[str] = None,
+        **kwargs,
+    ):
+        """Load from diffusers-format directory. Supports subfolder for single-component loading."""
+        path = Path(ensure_ldm_path(pretrained_model_name_or_path))
+        subfolder = kwargs.pop("subfolder", subfolder)
+
+        if subfolder in ("unet", "vae", "text_encoder", "local_adapter",
+                        "global_content_adapter", "global_text_adapter", "metadata_encoder"):
+            return load_component(path, subfolder)
+
+        if path.name in ("unet", "vae", "text_encoder", "local_adapter",
+                         "global_content_adapter", "global_text_adapter", "metadata_encoder"):
+            if (path / "config.json").exists():
+                ensure_ldm_path(path.parent)
+                return load_component(path.parent, path.name)
+
+        if not (path / "model_index.json").exists():
+            for _ in range(5):
+                parent = path.parent
+                if (parent / "model_index.json").exists():
+                    path = parent
+                    break
+                if parent == path:
+                    break
+                path = parent
+
+        components = load_components(path)
+        pipe = cls(
+            unet=components["unet"],
+            vae=components["vae"],
+            text_encoder=components["text_encoder"],
+            local_adapter=components["local_adapter"],
+            global_content_adapter=components["global_content_adapter"],
+            global_text_adapter=components["global_text_adapter"],
+            metadata_encoder=components["metadata_encoder"],
+            scheduler=components["scheduler"],
+            scale_factor=components["scale_factor"],
+        )
+        if device is not None:
+            pipe = pipe.to(device)
+        return pipe
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = "",
+        num_samples: int = 1,
+        height: int = 256,
+        width: int = 256,
+        num_inference_steps: int = 50,
+        eta: float = 0.0,
+        global_strength: float = 1.0,
+        text_strength: Optional[float] = None,
+        local_conditions: Optional[torch.Tensor] = None,
+        global_conditions: Optional[torch.Tensor] = None,
+        metadata: Optional[torch.Tensor] = None,
+        condition_resolution: int = 512,
+        guidance_scale: float = 1.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "numpy",
+        return_dict: bool = True,
+        save_memory: bool = False,
+    ):
+        device = next(self.crs_model.parameters()).device
+        if text_strength is None:
+            text_strength = global_strength
+
+        if isinstance(prompt, str):
+            prompts = [prompt] * num_samples
+        else:
+            prompts = list(prompt)
+            num_samples = len(prompts)
+
+        if save_memory:
+            self.crs_model.low_vram_shift(is_diffusing=False)
+
+        text_embedding = self.crs_model.get_learned_conditioning(prompts)
+
+        if local_conditions is None:
+            local_conditions = torch.zeros(
+                num_samples, 18, condition_resolution, condition_resolution,
+                device=device, dtype=torch.float32,
+            )
+        else:
+            local_conditions = local_conditions.to(device=device, dtype=torch.float32)
+
+        if global_conditions is None:
+            global_conditions = torch.zeros(
+                num_samples, 768, device=device, dtype=torch.float32,
+            )
+        else:
+            global_conditions = global_conditions.to(device=device, dtype=torch.float32)
+
+        if metadata is None:
+            metadata = torch.zeros(7, device=device, dtype=torch.float32)
+        else:
+            metadata = metadata.to(device=device, dtype=torch.float32)
+
+        cond = {
+            "local_control": [local_conditions],
+            "c_crossattn": [text_embedding],
+            "global_control": [global_conditions],
+            "metadata": [metadata],
+        }
+
+        do_cfg = guidance_scale > 1.0
+        if do_cfg:
+            if negative_prompt is None:
+                neg_prompts = [""] * num_samples
+            elif isinstance(negative_prompt, str):
+                neg_prompts = [negative_prompt] * num_samples
+            else:
+                neg_prompts = list(negative_prompt)
+            uc_text = self.crs_model.get_learned_conditioning(neg_prompts)
+            uncond = {
+                "local_control": [local_conditions],
+                "c_crossattn": [uc_text],
+                "global_control": [torch.zeros_like(global_conditions)],
+                "metadata": [metadata],
+            }
+
+        latent_shape = (num_samples, 4, height // 4, width // 4)
+        if latents is None:
+            latents = torch.randn(
+                latent_shape, device=device, generator=generator, dtype=torch.float32,
+            )
+        else:
+            latents = latents.to(device)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        if save_memory:
+            self.crs_model.low_vram_shift(is_diffusing=True)
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            t_batch = t.expand(num_samples)
+            if do_cfg:
+                noise_pred_cond = self.crs_model.apply_model(
+                    latents, t_batch, cond,
+                    metadata=metadata,
+                    global_strength=global_strength,
+                    text_strength=text_strength,
+                )
+                noise_pred_uncond = self.crs_model.apply_model(
+                    latents, t_batch, uncond,
+                    metadata=metadata,
+                    global_strength=global_strength,
+                    text_strength=text_strength,
+                )
+                noise_pred = noise_pred_uncond + guidance_scale * (
+                    noise_pred_cond - noise_pred_uncond
+                )
+            else:
+                noise_pred = self.crs_model.apply_model(
+                    latents, t_batch, cond,
+                    metadata=metadata,
+                    global_strength=global_strength,
+                    text_strength=text_strength,
+                )
+            latents = self.scheduler.step(
+                noise_pred, t, latents, eta=eta, generator=generator,
+            ).prev_sample
+
+        if output_type == "latent":
+            if not return_dict:
+                return (latents,)
+            return HSIGeneOutput(latents=latents)
+
+        if save_memory:
+            self.crs_model.low_vram_shift(is_diffusing=False)
+
+        images = self.crs_model.decode_first_stage(latents)
+        images = images.permute(0, 2, 3, 1).cpu().numpy()
+        images = images * 0.5 + 0.5
+        images = np.clip(images, 0, 1)
+
+        if not return_dict:
+            return (images,)
+        return HSIGeneOutput(images=images)

scheduler/scheduler_config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "DDIMScheduler",
+  "_diffusers_version": "0.37.0",
+  "beta_end": 0.02,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.0001,
+  "clip_sample": false,
+  "clip_sample_range": 1.0,
+  "dynamic_thresholding_ratio": 0.995,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "set_alpha_to_one": false,
+  "steps_offset": 0,
+  "thresholding": false,
+  "timestep_spacing": "leading",
+  "trained_betas": null
+}

text_encoder/__init__.py

@@ -0,0 +1 @@
+"""HSIGene text encoder component."""

text_encoder/config.json

@@ -0,0 +1,4 @@
+{
+  "_target": "hsigene.CLIPTextEncoder",
+  "version": "openai/clip-vit-large-patch14"
+}

text_encoder/model.py

@@ -0,0 +1,41 @@
+"""CLIP text encoder - same interface as FrozenCLIPEmbedder (forward(text) returns last_hidden_state)."""
+
+import torch
+import torch.nn as nn
+from transformers import CLIPTokenizer, CLIPTextModel
+
+
+class CLIPTextEncoder(nn.Module):
+    """CLIP text encoder wrapping transformers CLIPTokenizer + CLIPTextModel.
+    Same interface as FrozenCLIPEmbedder: forward(text) returns last_hidden_state.
+    """
+
+    def __init__(
+        self,
+        version: str = "openai/clip-vit-large-patch14",
+        max_length: int = 77,
+        freeze: bool = True,
+    ):
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.max_length = max_length
+        if freeze:
+            self.transformer.eval()
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, text):
+        """Encode text. Returns last_hidden_state (B, seq_len, dim)."""
+        if isinstance(text, str):
+            text = [text]
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            padding="max_length",
+            return_tensors="pt",
+        )
+        tokens = batch_encoding["input_ids"].to(next(self.parameters()).device)
+        outputs = self.transformer(input_ids=tokens)
+        return outputs.last_hidden_state

unet/__init__.py

@@ -0,0 +1,5 @@
+"""HSIGene UNet component."""
+
+from .model import HSIGeneUNet
+
+__all__ = ["HSIGeneUNet"]

unet/attention.py

@@ -0,0 +1,271 @@
+"""HSIGene attention modules - FeedForward, CrossAttention, SpatialTransformer."""
+
+from inspect import isfunction
+from typing import Optional, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import einsum
+
+from .utils import checkpoint, zero_module, exists
+
+try:
+    import xformers
+    import xformers.ops
+    XFORMERS_IS_AVAILABLE = True
+except ImportError:
+    XFORMERS_IS_AVAILABLE = False
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+import os
+_ATTN_PRECISION = os.environ.get("ATTN_PRECISION", "fp32")
+
+
+class GEGLU(nn.Module):
+    def __init__(self, dim_in, dim_out):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+            if not glu
+            else GEGLU(dim, inner_dim)
+        )
+        self.net = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+def Normalize(in_channels, num_groups=32):
+    return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x, context=None, mask=None):
+        h = self.heads
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
+        if _ATTN_PRECISION == "fp32":
+            with torch.autocast(enabled=False, device_type="cuda"):
+                q, k = q.float(), k.float()
+                sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
+        else:
+            sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
+        del q, k
+        if exists(mask):
+            mask = rearrange(mask, "b ... -> b (...)")
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, "b j -> (b h) () j", h=h)
+            sim.masked_fill_(~mask, max_neg_value)
+        sim = sim.softmax(dim=-1)
+        out = einsum("b i j, b j d -> b i d", sim, v)
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
+        return self.to_out(out)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = default(context_dim, query_dim)
+        self.heads = heads
+        self.dim_head = dim_head
+        self.scale = dim_head ** -0.5
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim),
+            nn.Dropout(dropout),
+        )
+        self.attention_op: Optional[Any] = None
+
+    def forward(self, x, context=None, mask=None):
+        q = self.to_q(x)
+        context = default(context, x)
+        k = self.to_k(context)
+        v = self.to_v(context)
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], self.heads, self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b * self.heads, t.shape[1], self.dim_head)
+            .contiguous(),
+            (q, k, v),
+        )
+        if XFORMERS_IS_AVAILABLE:
+            out = xformers.ops.memory_efficient_attention(
+                q, k, v, attn_bias=None, op=self.attention_op
+            )
+        else:
+            sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
+            sim = sim.softmax(dim=-1)
+            out = torch.einsum("b i j, b j d -> b i d", sim, v)
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, self.heads, out.shape[1], self.dim_head)
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], self.heads * self.dim_head)
+        )
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+    ATTENTION_MODES = {
+        "softmax": CrossAttention,
+        "softmax-xformers": MemoryEfficientCrossAttention,
+    }
+
+    def __init__(
+        self,
+        dim,
+        n_heads,
+        d_head,
+        dropout=0.0,
+        context_dim=None,
+        gated_ff=True,
+        checkpoint=True,
+        disable_self_attn=False,
+    ):
+        super().__init__()
+        attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILABLE else "softmax"
+        attn_cls = self.ATTENTION_MODES[attn_mode]
+        self.disable_self_attn = disable_self_attn
+        self.attn1 = attn_cls(
+            query_dim=dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+            context_dim=context_dim if self.disable_self_attn else None,
+        )
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = attn_cls(
+            query_dim=dim,
+            context_dim=context_dim,
+            heads=n_heads,
+            dim_head=d_head,
+            dropout=dropout,
+        )
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.norm3 = nn.LayerNorm(dim)
+        self.checkpoint = checkpoint
+
+    def forward(self, x, context=None):
+        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None):
+        x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x
+        x = self.attn2(self.norm2(x), context=context) + x
+        x = self.ff(self.norm3(x)) + x
+        return x
+
+
+class SpatialTransformer(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        n_heads,
+        d_head,
+        depth=1,
+        dropout=0.0,
+        context_dim=None,
+        disable_self_attn=False,
+        use_linear=False,
+        use_checkpoint=True,
+    ):
+        super().__init__()
+        if exists(context_dim) and not isinstance(context_dim, list):
+            context_dim = [context_dim]
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+        if not use_linear:
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
+        else:
+            self.proj_in = nn.Linear(in_channels, inner_dim)
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    n_heads,
+                    d_head,
+                    dropout=dropout,
+                    context_dim=context_dim[d] if isinstance(context_dim, list) else context_dim,
+                    disable_self_attn=disable_self_attn,
+                    checkpoint=use_checkpoint,
+                )
+                for d in range(depth)
+            ]
+        )
+        if not use_linear:
+            self.proj_out = zero_module(
+                nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+            )
+        else:
+            self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+        self.use_linear = use_linear
+
+    def forward(self, x, context=None):
+        if not isinstance(context, list):
+            context = [context]
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        if not self.use_linear:
+            x = self.proj_in(x)
+        x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+        if self.use_linear:
+            x = self.proj_in(x)
+        for i, block in enumerate(self.transformer_blocks):
+            ctx = context[i] if i < len(context) else context[0]
+            x = block(x, context=ctx)
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in

unet/config.json

@@ -0,0 +1,25 @@
+{
+  "_target": "hsigene.HSIGeneUNet",
+  "image_size": 32,
+  "in_channels": 4,
+  "model_channels": 320,
+  "out_channels": 4,
+  "num_res_blocks": 2,
+  "attention_resolutions": [
+    4,
+    2,
+    1
+  ],
+  "channel_mult": [
+    1,
+    2,
+    4,
+    4
+  ],
+  "use_checkpoint": true,
+  "num_heads": 8,
+  "use_spatial_transformer": true,
+  "transformer_depth": 1,
+  "context_dim": 768,
+  "legacy": false
+}

unet/diffusion.py

@@ -0,0 +1,608 @@
+"""HSIGene diffusion modules - UNet, ResBlock, etc. From openaimodel."""
+
+from abc import abstractmethod
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import (
+    checkpoint,
+    conv_nd,
+    linear,
+    zero_module,
+    normalization,
+    timestep_embedding,
+    exists,
+)
+from .attention import SpatialTransformer
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+    """Create a 1D, 2D, or 3D average pooling module."""
+    if dims == 1:
+        return nn.AvgPool1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.AvgPool2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.AvgPool3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def convert_module_to_f16(x):
+    pass
+
+
+def convert_module_to_f32(x):
+    pass
+
+
+class TimestepBlock(nn.Module):
+    """Any module where forward() takes timestep embeddings as a second argument."""
+
+    @abstractmethod
+    def forward(self, x, emb):
+        """Apply the module to `x` given `emb` timestep embeddings."""
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    """Sequential module that passes timestep embeddings to children that support it."""
+
+    def forward(self, x, emb, context=None):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            elif isinstance(layer, SpatialTransformer):
+                x = layer(x, context)
+            else:
+                x = layer(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """Upsampling layer with optional convolution."""
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        if use_conv:
+            self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        if self.dims == 3:
+            x = F.interpolate(
+                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+            )
+        else:
+            x = F.interpolate(x, scale_factor=2, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """Downsampling layer with optional convolution."""
+
+    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.dims = dims
+        stride = 2 if dims != 3 else (1, 2, 2)
+        if use_conv:
+            self.op = conv_nd(
+                dims, self.channels, self.out_channels, 3, stride=stride, padding=padding
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(TimestepBlock):
+    """Residual block with timestep conditioning."""
+
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        use_checkpoint=False,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_checkpoint = use_checkpoint
+        self.use_scale_shift_norm = use_scale_shift_norm
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            conv_nd(dims, channels, self.out_channels, 3, padding=1),
+        )
+
+        self.updown = up or down
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 3, padding=1)
+        else:
+            self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+    def forward(self, x, emb):
+        return checkpoint(self._forward, (x, emb), self.parameters(), self.use_checkpoint)
+
+    def _forward(self, x, emb):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = emb_out.chunk(2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AttentionBlock(nn.Module):
+    """Spatial self-attention block."""
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        use_checkpoint=False,
+        use_new_attention_order=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert channels % num_head_channels == 0
+            self.num_heads = channels // num_head_channels
+        self.use_checkpoint = use_checkpoint
+        self.norm = normalization(channels)
+        self.qkv = conv_nd(1, channels, channels * 3, 1)
+        self.attention = (
+            QKVAttention(self.num_heads)
+            if use_new_attention_order
+            else QKVAttentionLegacy(self.num_heads)
+        )
+        self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
+
+    def forward(self, x):
+        return checkpoint(self._forward, (x,), self.parameters(), True)
+
+    def _forward(self, x):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+class QKVAttentionLegacy(nn.Module):
+    """QKV attention - split heads before split qkv."""
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum("bct,bcs->bts", q * scale, k * scale)
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = torch.einsum("bts,bcs->bct", weight, v)
+        return a.reshape(bs, -1, length)
+
+
+class QKVAttention(nn.Module):
+    """QKV attention - split qkv before split heads."""
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.chunk(3, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum(
+            "bct,bcs->bts",
+            (q * scale).view(bs * self.n_heads, ch, length),
+            (k * scale).view(bs * self.n_heads, ch, length),
+        )
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = torch.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+        return a.reshape(bs, -1, length)
+
+
+class UNetModel(nn.Module):
+    """Full UNet with attention and timestep embedding."""
+
+    def __init__(
+        self,
+        image_size,
+        in_channels,
+        model_channels,
+        out_channels,
+        num_res_blocks,
+        attention_resolutions,
+        dropout=0,
+        channel_mult=(1, 2, 4, 8),
+        conv_resample=True,
+        dims=2,
+        num_classes=None,
+        use_checkpoint=False,
+        use_fp16=False,
+        num_heads=-1,
+        num_head_channels=-1,
+        num_heads_upsample=-1,
+        use_scale_shift_norm=False,
+        resblock_updown=False,
+        use_new_attention_order=False,
+        use_spatial_transformer=False,
+        transformer_depth=1,
+        context_dim=None,
+        n_embed=None,
+        legacy=True,
+        disable_self_attentions=None,
+        num_attention_blocks=None,
+        disable_middle_self_attn=False,
+        use_linear_in_transformer=False,
+    ):
+        super().__init__()
+        if use_spatial_transformer:
+            assert context_dim is not None
+        if context_dim is not None:
+            assert use_spatial_transformer
+            if hasattr(context_dim, "__iter__") and not isinstance(context_dim, (list, tuple)):
+                context_dim = list(context_dim)
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+        if num_heads == -1:
+            assert num_head_channels != -1
+        if num_head_channels == -1:
+            assert num_heads != -1
+
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        if isinstance(num_res_blocks, int):
+            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+        else:
+            assert len(num_res_blocks) == len(channel_mult)
+            self.num_res_blocks = num_res_blocks
+
+        self.attention_resolutions = attention_resolutions
+        self.dropout = dropout
+        self.channel_mult = channel_mult
+        self.conv_resample = conv_resample
+        self.num_classes = num_classes
+        self.use_checkpoint = use_checkpoint
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.num_heads = num_heads
+        self.num_head_channels = num_head_channels
+        self.num_heads_upsample = num_heads_upsample
+        self.predict_codebook_ids = n_embed is not None
+
+        time_embed_dim = model_channels * 4
+        self.time_embed = nn.Sequential(
+            linear(model_channels, time_embed_dim),
+            nn.SiLU(),
+            linear(time_embed_dim, time_embed_dim),
+        )
+
+        if num_classes is not None:
+            if isinstance(num_classes, int):
+                self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+            elif num_classes == "continuous":
+                self.label_emb = nn.Linear(1, time_embed_dim)
+            else:
+                raise ValueError()
+
+        self.input_blocks = nn.ModuleList(
+            [TimestepEmbedSequential(conv_nd(dims, in_channels, model_channels, 3, padding=1))]
+        )
+        self._feature_size = model_channels
+        input_block_chans = [model_channels]
+        ch = model_channels
+        ds = 1
+
+        for level, mult in enumerate(channel_mult):
+            for nr in range(self.num_res_blocks[level]):
+                layers = [
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=mult * model_channels,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = mult * model_channels
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads_cur = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+                    disabled_sa = (
+                        disable_self_attentions[level]
+                        if exists(disable_self_attentions)
+                        else False
+                    )
+                    if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
+                        attn_block = (
+                            AttentionBlock(
+                                ch,
+                                use_checkpoint=use_checkpoint,
+                                num_heads=num_heads,
+                                num_head_channels=dim_head,
+                                use_new_attention_order=use_new_attention_order,
+                            )
+                            if not use_spatial_transformer
+                            else SpatialTransformer(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                depth=transformer_depth,
+                                context_dim=context_dim,
+                                disable_self_attn=disabled_sa,
+                                use_linear=use_linear_in_transformer,
+                                use_checkpoint=use_checkpoint,
+                            )
+                        )
+                        layers.append(attn_block)
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                down_block = (
+                    ResBlock(
+                        ch,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=out_ch,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                        down=True,
+                    )
+                    if resblock_updown
+                    else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                )
+                self.input_blocks.append(TimestepEmbedSequential(down_block))
+                ch = out_ch
+                input_block_chans.append(ch)
+                ds *= 2
+                self._feature_size += ch
+
+        if num_head_channels == -1:
+            dim_head = ch // num_heads
+        else:
+            num_heads_cur = ch // num_head_channels
+            dim_head = num_head_channels
+        if legacy:
+            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+        mid_attn = (
+            AttentionBlock(
+                ch,
+                use_checkpoint=use_checkpoint,
+                num_heads=num_heads,
+                num_head_channels=dim_head,
+                use_new_attention_order=use_new_attention_order,
+            )
+            if not use_spatial_transformer
+            else SpatialTransformer(
+                ch,
+                num_heads,
+                dim_head,
+                depth=transformer_depth,
+                context_dim=context_dim,
+                disable_self_attn=disable_middle_self_attn,
+                use_linear=use_linear_in_transformer,
+                use_checkpoint=use_checkpoint,
+            )
+        )
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            mid_attn,
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+        )
+        self._feature_size += ch
+
+        self.output_blocks = nn.ModuleList([])
+        for level, mult in list(enumerate(channel_mult))[::-1]:
+            for i in range(self.num_res_blocks[level] + 1):
+                ich = input_block_chans.pop()
+                layers = [
+                    ResBlock(
+                        ch + ich,
+                        time_embed_dim,
+                        dropout,
+                        out_channels=model_channels * mult,
+                        dims=dims,
+                        use_checkpoint=use_checkpoint,
+                        use_scale_shift_norm=use_scale_shift_norm,
+                    )
+                ]
+                ch = model_channels * mult
+                if ds in attention_resolutions:
+                    if num_head_channels == -1:
+                        dim_head = ch // num_heads
+                    else:
+                        num_heads_cur = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        dim_head = (
+                            ch // num_heads if use_spatial_transformer else num_head_channels
+                        )
+                    disabled_sa = (
+                        disable_self_attentions[level]
+                        if exists(disable_self_attentions)
+                        else False
+                    )
+                    if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
+                        attn_block = (
+                            AttentionBlock(
+                                ch,
+                                use_checkpoint=use_checkpoint,
+                                num_heads=num_heads_upsample,
+                                num_head_channels=dim_head,
+                                use_new_attention_order=use_new_attention_order,
+                            )
+                            if not use_spatial_transformer
+                            else SpatialTransformer(
+                                ch,
+                                num_heads,
+                                dim_head,
+                                depth=transformer_depth,
+                                context_dim=context_dim,
+                                disable_self_attn=disabled_sa,
+                                use_linear=use_linear_in_transformer,
+                                use_checkpoint=use_checkpoint,
+                            )
+                        )
+                        layers.append(attn_block)
+                if level and i == self.num_res_blocks[level]:
+                    out_ch = ch
+                    up_block = (
+                        ResBlock(
+                            ch,
+                            time_embed_dim,
+                            dropout,
+                            out_channels=out_ch,
+                            dims=dims,
+                            use_checkpoint=use_checkpoint,
+                            use_scale_shift_norm=use_scale_shift_norm,
+                            up=True,
+                        )
+                        if resblock_updown
+                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+                    )
+                    layers.append(up_block)
+                    ds //= 2
+                self.output_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+
+        self.out = nn.Sequential(
+            normalization(ch),
+            nn.SiLU(),
+            zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+        )
+        if self.predict_codebook_ids:
+            self.id_predictor = nn.Sequential(
+                normalization(ch),
+                conv_nd(dims, model_channels, n_embed, 1),
+            )
+
+    def forward(self, x, timesteps=None, metadata=None, context=None, y=None, **kwargs):
+        assert (y is not None) == (self.num_classes is not None)
+        hs = []
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        emb = self.time_embed(t_emb)
+        if metadata is not None:
+            if isinstance(metadata, (list, tuple)) and len(metadata) == 1:
+                metadata = metadata[0]
+            emb = emb + metadata
+
+        if self.num_classes is not None:
+            assert y.shape[0] == x.shape[0]
+            emb = emb + self.label_emb(y)
+
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb, context)
+            hs.append(h)
+        h = self.middle_block(h, emb, context)
+        for module in self.output_blocks:
+            h = torch.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context)
+        h = h.type(x.dtype)
+        if self.predict_codebook_ids:
+            return self.id_predictor(h)
+        return self.out(h)

unet/model.py

@@ -0,0 +1,35 @@
+"""HSIGene UNet - LocalControlUNetModel for hyperspectral generation."""
+
+import torch
+
+from .diffusion import UNetModel
+from .utils import timestep_embedding
+
+
+class HSIGeneUNet(UNetModel):
+    """UNet that accepts metadata and local_control from LocalAdapter."""
+
+    def forward(
+        self,
+        x,
+        timesteps=None,
+        metadata=None,
+        context=None,
+        local_control=None,
+        meta=False,
+        **kwargs,
+    ):
+        hs = []
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        emb = self.time_embed(t_emb) + metadata
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            h = module(h, emb, context)
+            hs.append(h)
+        h = self.middle_block(h, emb, context)
+        h += local_control.pop()
+        for module in self.output_blocks:
+            h = torch.cat([h, hs.pop() + local_control.pop()], dim=1)
+            h = module(h, emb, context)
+        h = h.type(x.dtype)
+        return self.out(h)

unet/utils.py

@@ -0,0 +1,90 @@
+"""HSIGene utilities - no ldm/models imports."""
+
+import math
+import torch
+import torch.nn as nn
+from einops import repeat
+
+
+def exists(val):
+    return val is not None
+
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    if repeat_only:
+        return repeat(timesteps, "b -> b d", d=dim)
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+def conv_nd(dims, *args, **kwargs):
+    if dims == 1:
+        return nn.Conv1d(*args, **kwargs)
+    elif dims == 2:
+        return nn.Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return nn.Conv3d(*args, **kwargs)
+    raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+    return nn.Linear(*args, **kwargs)
+
+
+def zero_module(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def checkpoint(func, inputs, params, flag):
+    if flag:
+        return _CheckpointFunction.apply(func, len(inputs), *(tuple(inputs) + tuple(params)))
+    return func(*inputs)
+
+
+class _CheckpointFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, run_function, length, *args):
+        ctx.run_function = run_function
+        ctx.input_tensors = list(args[:length])
+        ctx.input_params = list(args[length:])
+        ctx.gpu_autocast_kwargs = {
+            "enabled": torch.is_autocast_enabled(),
+            "dtype": torch.get_autocast_gpu_dtype(),
+            "cache_enabled": torch.is_autocast_cache_enabled(),
+        }
+        with torch.no_grad():
+            output_tensors = ctx.run_function(*ctx.input_tensors)
+        return output_tensors
+
+    @staticmethod
+    def backward(ctx, *output_grads):
+        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+        with torch.enable_grad(), torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
+            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+            output_tensors = ctx.run_function(*shallow_copies)
+        input_grads = torch.autograd.grad(
+            output_tensors,
+            ctx.input_tensors + ctx.input_params,
+            output_grads,
+            allow_unused=True,
+        )
+        return (None, None) + input_grads
+
+
+def normalization(channels):
+    return GroupNorm32(32, channels)
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+

vae/__init__.py

@@ -0,0 +1 @@
+"""HSIGene VAE component."""

vae/config.json

@@ -0,0 +1,21 @@
+{
+  "_target": "hsigene.HSIGeneAutoencoderKL",
+  "in_channels": 48,
+  "out_channels": 48,
+  "latent_channels": 96,
+  "embed_dim": 4,
+  "block_out_channels": [
+    64,
+    128,
+    256
+  ],
+  "num_res_blocks": 4,
+  "attn_resolutions": [
+    16,
+    32,
+    64
+  ],
+  "dropout": 0.0,
+  "double_z": true,
+  "resolution": 256
+}

vae/model.py

@@ -0,0 +1,90 @@
+"""HSIGene AutoencoderKL - nn.Module, no Lightning. Loss = Identity."""
+
+import torch
+import torch.nn as nn
+
+from .vae_blocks import Encoder, Decoder, DiagonalGaussianDistribution
+
+
+class AutoencoderKL(nn.Module):
+    """
+    AutoencoderKL - nn.Module (not Lightning).
+    Uses Encoder, Decoder, quant_conv, post_quant_conv.
+    encode() returns posterior, decode() takes z.
+    Loss = Identity (no-op).
+    """
+
+    def __init__(
+        self,
+        ddconfig,
+        embed_dim=4,
+        lossconfig=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.encoder = Encoder(**ddconfig)
+        self.decoder = Decoder(**ddconfig)
+        assert ddconfig.get("double_z", True)
+        z_channels = ddconfig["z_channels"]
+        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * embed_dim, 1)
+        self.post_quant_conv = nn.Conv2d(embed_dim, z_channels, 1)
+        self.embed_dim = embed_dim
+        self.loss = nn.Identity()
+
+    def encode(self, x):
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(moments, deterministic=True)
+        return posterior
+
+    def decode(self, z):
+        z = self.post_quant_conv(z)
+        return self.decoder(z)
+
+    def forward(self, input, sample_posterior=True):
+        posterior = self.encode(input)
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec, posterior
+
+
+class HSIGeneAutoencoderKL(AutoencoderKL):
+    """
+    HSIGene VAE with diffusers-style config.
+    Accepts in_channels, out_channels, latent_channels, block_out_channels.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 48,
+        out_channels: int = 48,
+        latent_channels: int = 96,
+        embed_dim: int = 4,
+        block_out_channels: tuple = (64, 128, 256),
+        num_res_blocks: int = 4,
+        attn_resolutions: tuple = (16, 32, 64),
+        dropout: float = 0.0,
+        double_z: bool = True,
+        resolution: int = 256,
+        **kwargs,
+    ):
+        ch = block_out_channels[0]
+        ch_mult = tuple(
+            block_out_channels[i] // ch for i in range(len(block_out_channels))
+        )
+        ddconfig = dict(
+            double_z=double_z,
+            z_channels=latent_channels,
+            resolution=resolution,
+            in_channels=in_channels,
+            out_ch=out_channels,
+            ch=ch,
+            ch_mult=list(ch_mult),
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=list(attn_resolutions),
+            dropout=dropout,
+        )
+        super().__init__(ddconfig=ddconfig, embed_dim=embed_dim, **kwargs)

vae/utils.py

@@ -0,0 +1,10 @@
+"""VAE utilities."""
+
+import torch.nn as nn
+
+
+def zero_module(module):
+    """Zero out the parameters of a module and return it."""
+    for p in module.parameters():
+        p.detach().zero_()
+    return module

vae/vae_blocks.py

@@ -0,0 +1,441 @@
+"""HSIGene VAE blocks - ResnetBlock, Encoder, Decoder, DiagonalGaussianDistribution."""
+
+from typing import Optional, Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+try:
+    import xformers
+    import xformers.ops
+    XFORMERS_IS_AVAILABLE = True
+except ImportError:
+    XFORMERS_IS_AVAILABLE = False
+
+
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels, num_groups=32):
+    return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout,
+        temb_channels=512,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if temb_channels > 0:
+            self.temb_proj = nn.Linear(temb_channels, out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = nn.Conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = nn.Conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w).permute(0, 2, 1)
+        k = k.reshape(b, c, h * w)
+        w_ = torch.bmm(q, k) * (int(c) ** -0.5)
+        w_ = F.softmax(w_, dim=2)
+        v = v.reshape(b, c, h * w)
+        h_ = torch.bmm(v, w_.permute(0, 2, 1))
+        h_ = h_.reshape(b, c, h, w)
+        h_ = self.proj_out(h_)
+        return x + h_
+
+
+class MemoryEfficientAttnBlock(nn.Module):
+    """AttnBlock using xformers when available."""
+
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.attention_op: Optional[Any] = None
+
+    def forward(self, x):
+        h_ = self.norm(x)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+        B, C, H, W = q.shape
+        q, k, v = map(lambda t: rearrange(t, "b c h w -> b (h w) c"), (q, k, v))
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(B, t.shape[1], 1, C)
+            .permute(0, 2, 1, 3)
+            .reshape(B * 1, t.shape[1], C)
+            .contiguous(),
+            (q, k, v),
+        )
+        out = xformers.ops.memory_efficient_attention(
+            q, k, v, attn_bias=None, op=self.attention_op
+        )
+        out = (
+            out.unsqueeze(0)
+            .reshape(B, 1, out.shape[1], C)
+            .permute(0, 2, 1, 3)
+            .reshape(B, out.shape[1], C)
+        )
+        out = rearrange(out, "b (h w) c -> b c h w", b=B, h=H, w=W, c=C)
+        out = self.proj_out(out)
+        return x + out
+
+
+def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
+    assert attn_type in ["vanilla", "vanilla-xformers", "none"]
+    if XFORMERS_IS_AVAILABLE and attn_type == "vanilla":
+        attn_type = "vanilla-xformers"
+    if attn_type == "vanilla":
+        return AttnBlock(in_channels)
+    elif attn_type == "vanilla-xformers":
+        return MemoryEfficientAttnBlock(in_channels)
+    elif attn_type == "none":
+        return nn.Identity()
+    raise NotImplementedError(f"attn_type {attn_type}")
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        ch,
+        out_ch,
+        ch_mult=(1, 2, 4, 8),
+        num_res_blocks,
+        attn_resolutions,
+        dropout=0.0,
+        resamp_with_conv=True,
+        in_channels,
+        resolution,
+        z_channels,
+        double_z=True,
+        use_linear_attn=False,
+        attn_type="vanilla",
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        if use_linear_attn:
+            attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(num_res_blocks):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.norm_out = Normalize(block_in)
+        self.conv_out = nn.Conv2d(
+            block_in, 2 * z_channels if double_z else z_channels,
+            kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, x):
+        temb = None
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        ch,
+        out_ch,
+        ch_mult=(1, 2, 4, 8),
+        num_res_blocks,
+        attn_resolutions,
+        dropout=0.0,
+        resamp_with_conv=True,
+        in_channels,
+        resolution,
+        z_channels,
+        give_pre_end=False,
+        tanh_out=False,
+        use_linear_attn=False,
+        attn_type="vanilla",
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        if use_linear_attn:
+            attn_type = "linear"
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+        self.tanh_out = tanh_out
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+        self.mid.block_2 = ResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+        )
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    ResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(make_attn(block_in, attn_type=attn_type))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)
+
+        self.norm_out = Normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z):
+        self.last_z_shape = z.shape
+        temb = None
+        h = self.conv_in(z)
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+        if self.give_pre_end:
+            return h
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        if self.tanh_out:
+            h = torch.tanh(h)
+        return h
+
+
+class DiagonalGaussianDistribution:
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -20.0, 0.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean, device=parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(
+            self.mean.shape, device=self.parameters.device
+        )
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.tensor(0.0, device=self.parameters.device)
+        if other is None:
+            return 0.5 * torch.sum(
+                torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                dim=[1, 2, 3],
+            )
+        return 0.5 * torch.sum(
+            torch.pow(self.mean - other.mean, 2) / other.var
+            + self.var / other.var
+            - 1.0
+            - self.logvar
+            + other.logvar,
+            dim=[1, 2, 3],
+        )
+
+    def nll(self, sample, dims=[1, 2, 3]):
+        if self.deterministic:
+            return torch.tensor(0.0, device=self.parameters.device)
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self):
+        return self.mean