Add files using upload-large-folder tool

README.md

@@ -0,0 +1,88 @@
+---
+license: apache-2.0
+library_name: diffusers
+tags:
+  - aerogen
+  - remote-sensing
+  - object-detection
+  - latent-diffusion
+  - bounding-box
+  - arxiv:2411.15497
+pipeline_tag: text-to-image
+language:
+  - en
+---
+
+# BiliSakura/AeroGen
+
+**Aerial image generation** conditioned on bounding boxes (horizontal or rotated) and object categories. AeroGen is the first model to simultaneously support horizontal and rotated bounding box condition generation for remote sensing imagery.
+
+Converted to diffusers format. **Self-contained** — no external code repo needed; all required code is bundled.
+
+## Model Details
+
+- **Model type**: Latent diffusion with UNet + VAE + CLIP text encoder + RBoxEncoder (condition encoder)
+- **Conditioning**: Bounding boxes (8 coords for rotated, 4 for axis-aligned), category CLIP embeddings, spatial masks
+- **Scheduler**: DDIMScheduler, 1000 steps, scaled_linear
+- **Output**: 512×512 RGB aerial images
+- **License**: Apache 2.0
+
+### Repository Structure
+
+| Component         | Path                 |
+|-------------------|----------------------|
+| Pipeline          | `pipeline.py`        |
+| UNet              | `unet/`              |
+| VAE               | `vae/`               |
+| Text encoder      | `text_encoder/`      |
+| Condition encoder | `condition_encoder/` |
+| Scheduler         | `scheduler/`         |
+| Config            | `model_index.json`   |
+
+## Inference
+
+**Dependencies:** `pip install diffusers transformers torch einops safetensors pyyaml`
+
+```python
+from diffusers import DiffusionPipeline
+import torch
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/AeroGen",
+    trust_remote_code=True,
+)
+pipe = pipe.to("cuda")
+```
+
+### Conditioning Format
+
+| Input               | Shape        | Description                                              |
+|---------------------|--------------|----------------------------------------------------------|
+| `bboxes`            | (B, N, 8)    | Rotated box corners [x1,y1,x2,y2,x3,y3,x4,y4], normalized |
+| `bboxes`            | (B, N, 4)    | Axis-aligned [x1,y1,x2,y2], normalized                   |
+| `category_conditions` | (B, N, 768) | CLIP text embeddings per object (e.g. encode class name) |
+| `mask_conditions`   | (B, N, 64, 64) | Spatial mask per object (64×64 for 512px output)      |
+| `mask_vector`       | (B, N)       | 1 = valid object, 0 = padding                            |
+
+For layout preparation and DIOR-R format, see the [original AeroGen repo](https://github.com/Sonettoo/AeroGen).
+
+## Model Sources
+
+- **Source**: [Sonetto702/AeroGen](https://huggingface.co/Sonetto702/AeroGen)
+- **Paper**: [AeroGen: Enhancing Remote Sensing Object Detection with Diffusion-Driven Data Generation](https://arxiv.org/abs/2411.15497)
+- **Original repo**: [Sonettoo/AeroGen](https://github.com/Sonettoo/AeroGen)
+- **Conversion**: Checkpoint converted to diffusers format (self-contained, no external repo)
+
+## Citation
+
+```bibtex
+@inproceedings{tangAeroGenEnhancingRemote2025,
+  title = {{{AeroGen}}: {{Enhancing Remote Sensing Object Detection}} with {{Diffusion-Driven Data Generation}}},
+  shorttitle = {{{AeroGen}}},
+  booktitle = {{{CVPR}}},
+  author = {Tang, Datao and Cao, Xiangyong and Wu, Xuan and Li, Jialin and Yao, Jing and Bai, Xueru and Jiang, Dongsheng and Li, Yin and Meng, Deyu},
+  year = 2025,
+  pages = {3614--3624},
+  urldate = {2025-11-20}
+}
+```

condition_encoder/__init__.py

@@ -0,0 +1,5 @@
+"""AeroGen condition encoder (RBoxEncoder)."""
+
+from .rbox_encoder import RBoxEncoder
+
+__all__ = ["RBoxEncoder"]

condition_encoder/config.json

@@ -0,0 +1,7 @@
+{
+  "target": "condition_encoder.rbox_encoder.RBoxEncoder",
+  "params": {
+    "in_dim": 768,
+    "out_dim": 768
+  }
+}

condition_encoder/diffusion_pytorch_model.safetensors

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

condition_encoder/rbox_encoder.py

@@ -0,0 +1,67 @@
+"""
+RBoxEncoder - pure PyTorch, no ldm/bldm dependency.
+
+Encodes rotated bounding boxes (8 coords) with Fourier embedding and text embeddings.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class FourierEmbedder:
+    def __init__(self, num_freqs=64, temperature=100):
+        self.num_freqs = num_freqs
+        self.temperature = temperature
+        self.freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
+
+    @torch.no_grad()
+    def __call__(self, x, cat_dim=-1):
+        out = []
+        for freq in self.freq_bands:
+            out.append(torch.sin(freq * x))
+            out.append(torch.cos(freq * x))
+        return torch.cat(out, cat_dim)
+
+
+class RBoxEncoder(nn.Module):
+    """Encoder for rotated bounding boxes (8 coords) with text embeddings."""
+
+    def __init__(self, in_dim, out_dim, fourier_freqs=8):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
+        self.position_dim = fourier_freqs * 2 * 8  # 2 is sin&cos, 8 is xyxyxyxy
+
+        self.linears = nn.Sequential(
+            nn.Linear(self.in_dim + self.position_dim, 512),
+            nn.SiLU(),
+            nn.Linear(512, 512),
+            nn.SiLU(),
+            nn.Linear(512, out_dim),
+        )
+
+        self.null_text_feature = nn.Parameter(torch.zeros([self.in_dim]))
+        self.null_position_feature = nn.Parameter(torch.zeros([self.position_dim]))
+
+    def forward(self, boxes=None, masks=None, text_embeddings=None, **kwargs):
+        # Pipeline passes boxes=[bboxes], masks=[mask_vector], text_embeddings=[category_conditions]
+        boxes = (boxes or kwargs.get("boxes", [[]]))[0]
+        masks = (masks or kwargs.get("masks", [[]]))[0]
+        text_embeddings = (text_embeddings or kwargs.get("text_embeddings", [[]]))[0]
+
+        B, N, _ = boxes.shape
+        masks = masks.unsqueeze(-1)
+
+        xyxy_embedding = self.fourier_embedder(boxes)  # B*N*8 --> B*N*C
+
+        text_null = self.null_text_feature.view(1, 1, -1)
+        xyxy_null = self.null_position_feature.view(1, 1, -1)
+
+        text_embeddings = text_embeddings * masks + (1 - masks) * text_null
+        xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null
+
+        objs = self.linears(torch.cat([text_embeddings, xyxy_embedding], dim=-1))
+        assert objs.shape == torch.Size([B, N, self.out_dim])
+        return objs

model_index.json

@@ -0,0 +1,25 @@
+{
+  "_class_name": ["pipeline", "AeroGenPipeline"],
+  "_diffusers_version": "0.25.0",
+  "condition_encoder": [
+    "pipeline",
+    "AeroGenPipeline"
+  ],
+  "scheduler": [
+    "diffusers",
+    "DDIMScheduler"
+  ],
+  "text_encoder": [
+    "pipeline",
+    "AeroGenPipeline"
+  ],
+  "unet": [
+    "pipeline",
+    "AeroGenPipeline"
+  ],
+  "vae": [
+    "pipeline",
+    "AeroGenPipeline"
+  ],
+  "scale_factor": 0.18215
+}

modular_pipeline.py

@@ -0,0 +1,167 @@
+"""
+AeroGen modular components: scheduler config, component loading, and path setup.
+
+Self-contained - no ldm/bldm. Scheduler is created in-code (no scheduler/ folder required).
+"""
+
+import importlib
+import json
+import sys
+from pathlib import Path
+from typing import Optional, Union
+
+from diffusers import DDIMScheduler
+
+# Ensure model dir is on path for local module imports (unet, text_encoder, condition_encoder)
+_pipeline_dir = Path(__file__).resolve().parent
+if str(_pipeline_dir) not in sys.path:
+    sys.path.insert(0, str(_pipeline_dir))
+
+# Default DDIM scheduler config (matches scheduler/scheduler_config.json)
+DEFAULT_SCHEDULER_CONFIG = {
+    "num_train_timesteps": 1000,
+    "beta_start": 0.00085,
+    "beta_end": 0.012,
+    "beta_schedule": "scaled_linear",
+    "clip_sample": False,
+    "set_alpha_to_one": False,
+    "prediction_type": "epsilon",
+}
+
+
+def ensure_ldm_path(pretrained_model_name_or_path: Union[str, Path]) -> Path:
+    """Add model repo to path so local modules 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 ensure_ldm_path_from_config(config_path: str) -> None:
+    """Walk up from config file dir to find project root and add to path."""
+    d = Path(config_path).resolve().parent
+    for _ in range(10):
+        if (d / "pipeline.py").exists() or (d / "unet").is_dir():
+            s = str(d)
+            if s not in sys.path:
+                sys.path.insert(0, s)
+            return
+        parent = d.parent
+        if parent == d:
+            break
+        d = parent
+
+
+def _get_class_from_string(target: str):
+    """Resolve class from dotted path (diffusers-style, no OmegaConf)."""
+    module_path, cls_name = target.rsplit(".", 1)
+    mod = importlib.import_module(module_path)
+    return getattr(mod, cls_name)
+
+
+def _instantiate_from_config(config: dict):
+    """Instantiate from dict with 'target' and 'params' (diffusers-style, no OmegaConf)."""
+    if not isinstance(config, dict) or "target" not in config:
+        raise KeyError("Expected key 'target' to instantiate.")
+    cls = _get_class_from_string(config["target"])
+    params = dict(config.get("params") or {})
+    params.pop("ckpt_path", None)
+    params.pop("ignore_keys", None)
+    params.pop("target", None)  # avoid passing target into constructor
+    return cls(**params)
+
+
+def create_scheduler(model_path: Path) -> DDIMScheduler:
+    """Create DDIMScheduler from path/scheduler if exists, else from defaults."""
+    scheduler_path = model_path / "scheduler"
+    if scheduler_path.exists() and (scheduler_path / "scheduler_config.json").exists():
+        return DDIMScheduler.from_pretrained(scheduler_path)
+    return DDIMScheduler(**DEFAULT_SCHEDULER_CONFIG)
+
+
+def load_component(model_path: Path, name: str):
+    """Load a custom component (unet, vae, text_encoder, condition_encoder).
+
+    VAE: Uses diffusers AutoencoderKL.from_pretrained when saved in diffusers format
+         (config has down_block_types, no target). Otherwise uses target/params.
+
+    When diffusers loads a single component, it passes the component subfolder path
+    (e.g. .../unet). We detect that and use it directly.
+    """
+    import torch
+    path = Path(model_path)
+    # Ensure model root is on sys.path for imports (unet, text_encoder, condition_encoder)
+    root = path.parent if path.name in ("unet", "vae", "text_encoder", "condition_encoder") and (path / "config.json").exists() else path
+    ensure_ldm_path(root)
+    # If path is already a component folder (has config.json), use it directly
+    if (path / "config.json").exists() and path.name in ("unet", "vae", "text_encoder", "condition_encoder"):
+        comp_path = path
+    else:
+        comp_path = path / name
+    with open(comp_path / "config.json") as f:
+        cfg = json.load(f)
+
+    # Diffusers native format (e.g. AutoencoderKL.save_pretrained): no "target" key
+    if "target" not in cfg and name == "vae":
+        from diffusers import AutoencoderKL
+        return AutoencoderKL.from_pretrained(comp_path)
+
+    component = _instantiate_from_config(cfg)
+    safetensors_path = comp_path / "diffusion_pytorch_model.safetensors"
+    bin_path = comp_path / "diffusion_pytorch_model.bin"
+    if safetensors_path.exists():
+        import safetensors.torch
+        state = safetensors.torch.load_file(str(safetensors_path))
+    elif bin_path.exists():
+        try:
+            state = torch.load(str(bin_path), map_location="cpu", weights_only=True)
+        except TypeError:
+            state = torch.load(str(bin_path), map_location="cpu")
+    else:
+        raise FileNotFoundError(
+            f"No weights in {comp_path} "
+            "(expected diffusion_pytorch_model.safetensors or .bin)"
+        )
+    # UNet: AeroGenUNet2DConditionModel wraps UNetModel in self.model, so expects "model.xxx" keys.
+    # Older checkpoints may have been saved without the "model." prefix.
+    if name == "unet" and state and not any(k.startswith("model.") for k in state.keys()):
+        state = {"model." + k: v for k, v in state.items()}
+    component.load_state_dict(state, strict=True)
+    component.eval()
+    return component
+
+
+def load_components(
+    model_path: Union[str, Path],
+) -> dict:
+    """Load all pipeline components. Returns dict with unet, vae, text_encoder, condition_encoder, scheduler, scale_factor."""
+    path = Path(ensure_ldm_path(model_path))
+    # If path points to a component subfolder (e.g. .../unet), use parent as model root
+    if path.name in ("unet", "vae", "text_encoder", "condition_encoder") and (path / "config.json").exists():
+        path = path.parent
+    scheduler = create_scheduler(path)
+    unet = load_component(path, "unet")
+    vae = load_component(path, "vae")
+    text_encoder = load_component(path, "text_encoder")
+    condition_encoder = load_component(path, "condition_encoder")
+
+    scale_factor = 0.18215
+    model_index_path = path / "model_index.json"
+    if model_index_path.exists():
+        with open(model_index_path) as f:
+            model_index = json.load(f)
+        scale_factor = model_index.get("scale_factor", scale_factor)
+
+    return {
+        "unet": unet,
+        "vae": vae,
+        "text_encoder": text_encoder,
+        "condition_encoder": condition_encoder,
+        "scheduler": scheduler,
+        "scale_factor": scale_factor,
+    }

pipeline.py

@@ -0,0 +1,377 @@
+"""
+AeroGen Pipeline using native HuggingFace Diffusers.
+
+This module provides a DiffusionPipeline subclass that wraps AeroGen's
+custom UNet, condition encoder, VAE, and text encoder into a standard
+diffusers pipeline interface, using DDIMScheduler for the denoising loop.
+
+Usage:
+    # Load from config + checkpoint
+    pipeline = AeroGenPipeline.from_pretrained_checkpoint(
+        config_path="configs/.../v1-finetune-DIOR-R.yaml",
+        checkpoint_path="./ckpt/aerogen_diorr_last.ckpt",
+    )
+
+    # Load from diffusers-format (after convert_to_diffusers.py)
+    pipeline = AeroGenPipeline.from_pretrained("/path/to/AeroGen")
+"""
+
+import json
+import os
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Union
+
+# Ensure model repo is on path for trust_remote_code / custom_pipeline loading
+_pipeline_dir = Path(__file__).resolve().parent
+if str(_pipeline_dir) not in sys.path:
+    sys.path.insert(0, str(_pipeline_dir))
+
+import einops
+import numpy as np
+import torch
+import yaml
+from diffusers import DDIMScheduler, DiffusionPipeline
+from diffusers.utils import BaseOutput
+from PIL import Image
+
+from modular_pipeline import (
+    ensure_ldm_path,
+    ensure_ldm_path_from_config,
+    load_component,
+    load_components,
+    create_scheduler,
+    _instantiate_from_config,
+)
+
+
+@dataclass
+class AeroGenPipelineOutput(BaseOutput):
+    """Output class for AeroGen pipeline.
+
+    Attributes:
+        images: List of generated PIL images.
+    """
+
+    images: List[Image.Image]
+
+
+class AeroGenPipeline(DiffusionPipeline):
+    """Pipeline for AeroGen: conditional aerial image generation with
+    bounding box and category controls.
+
+    This pipeline wraps AeroGen's custom components (UNet, condition encoder,
+    VAE, text encoder) and uses a native diffusers DDIMScheduler for the
+    denoising loop, replacing the original custom DDIM sampler.
+
+    Args:
+        unet: The custom UNet model (openaimodel_bbox_v2.UNetModel).
+        scheduler: A diffusers DDIMScheduler instance.
+        vae: The VAE model (AutoencoderKL) for latent encoding/decoding.
+        text_encoder: The frozen CLIP text encoder for prompt conditioning.
+        condition_encoder: The RBoxEncoder or BoxEncoder for bbox conditioning.
+        scale_factor: VAE latent scale factor (default: 0.18215 for SD 1.x).
+    """
+
+    def __init__(
+        self,
+        unet: torch.nn.Module,
+        scheduler: DDIMScheduler,
+        vae: torch.nn.Module,
+        text_encoder: torch.nn.Module,
+        condition_encoder: torch.nn.Module,
+        scale_factor: float = 0.18215,
+    ):
+        super().__init__()
+        self.register_modules(
+            unet=unet,
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            condition_encoder=condition_encoder,
+        )
+        self.vae_scale_factor = scale_factor
+
+    @property
+    def device(self) -> torch.device:
+        """Return the device of the pipeline's first nn.Module parameter."""
+        for module in [self.unet, self.vae, self.text_encoder, self.condition_encoder]:
+            if isinstance(module, torch.nn.Module):
+                params = list(module.parameters())
+                if params:
+                    return params[0].device
+        return torch.device("cpu")
+
+    @property
+    def _execution_device(self) -> torch.device:
+        return self.device
+
+    @classmethod
+    def from_pretrained_checkpoint(
+        cls,
+        config_path: str,
+        checkpoint_path: str,
+        device: str = "cuda",
+    ) -> "AeroGenPipeline":
+        """Load an AeroGenPipeline from a YAML config and checkpoint.
+
+        DEPRECATED: ldm/bldm have been removed. Use from_pretrained() with a
+        diffusers-format model (converted via convert_to_diffusers_lowvram.py).
+        """
+        raise NotImplementedError(
+            "from_pretrained_checkpoint is no longer supported (ldm/bldm removed). "
+            "Use AeroGenPipeline.from_pretrained() with a diffusers-format 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,
+    ) -> Union["AeroGenPipeline", torch.nn.Module]:
+        """Load AeroGenPipeline from a diffusers-format directory.
+
+        Supports native diffusers loading via DiffusionPipeline.from_pretrained(..., trust_remote_code=True).
+        When subfolder is provided (e.g. by diffusers for component loading), returns only that component.
+
+        Args:
+            pretrained_model_name_or_path: Path to the diffusers-format
+                directory or HuggingFace repo ID.
+            device: Device to load the model onto.
+            subfolder: If set, load only this component (unet, vae, text_encoder, condition_encoder).
+
+        Returns:
+            An AeroGenPipeline instance, or a single component if subfolder is set.
+        """
+        path = Path(ensure_ldm_path(pretrained_model_name_or_path))
+
+        # Single-component loading (for diffusers trust_remote_code component loading)
+        subfolder = kwargs.pop("subfolder", subfolder)
+        if subfolder in ("unet", "vae", "text_encoder", "condition_encoder"):
+            return load_component(path, subfolder)
+
+        # When diffusers loads a component, it passes the component subfolder path directly
+        if path.name in ("unet", "vae", "text_encoder", "condition_encoder") and (path / "config.json").exists():
+            ensure_ldm_path(path.parent)  # Ensure model root is on sys.path for imports
+            return load_component(path.parent, path.name)
+
+        # Ensure we have model root (diffusers may pass a subfolder when loading full pipeline)
+        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"],
+            scheduler=components["scheduler"],
+            vae=components["vae"],
+            text_encoder=components["text_encoder"],
+            condition_encoder=components["condition_encoder"],
+            scale_factor=components["scale_factor"],
+        )
+
+        if device is not None:
+            pipe = pipe.to(device)
+        return pipe
+
+    def _encode_prompt(self, prompt: Union[str, List[str]]) -> torch.Tensor:
+        """Encode text prompt(s) using the frozen CLIP text encoder."""
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        return self.text_encoder.encode(prompt)
+
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """Decode latent representations using the VAE."""
+        latents = (1.0 / self.vae_scale_factor) * latents
+        image = self.vae.decode(latents)
+        return image
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        bboxes: torch.Tensor,
+        category_conditions: torch.Tensor,
+        mask_conditions: torch.Tensor,
+        mask_vector: torch.Tensor,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        eta: float = 0.2,
+        height: int = 512,
+        width: int = 512,
+        num_images_per_prompt: int = 1,
+        generator: Optional[torch.Generator] = None,
+        output_type: str = "pil",
+    ) -> AeroGenPipelineOutput:
+        """Generate aerial images conditioned on bounding boxes and categories.
+
+        Args:
+            prompt: Text prompt(s) describing the aerial scene.
+            bboxes: Bounding box coordinates tensor of shape (B, N, 8) for
+                rotated boxes or (B, N, 4) for axis-aligned boxes.
+            category_conditions: Category embedding tensor of shape
+                (B, N, 768).
+            mask_conditions: Spatial mask tensor of shape (B, N, H, W).
+            mask_vector: Binary vector indicating valid objects, shape (B, N).
+            num_inference_steps: Number of DDIM denoising steps.
+            guidance_scale: Classifier-free guidance scale. Values > 1.0
+                enable guidance.
+            eta: DDIM eta parameter controlling stochasticity.
+            height: Output image height (must be divisible by 8).
+            width: Output image width (must be divisible by 8).
+            num_images_per_prompt: Number of images to generate per prompt.
+            generator: Optional torch.Generator for reproducibility.
+            output_type: Output format, either "pil" for PIL images or
+                "tensor" for raw image tensors.
+
+        Returns:
+            AeroGenPipelineOutput with the generated images.
+        """
+        device = self._execution_device
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        batch_size = len(prompt)
+
+        # Repeat conditions for num_images_per_prompt
+        if num_images_per_prompt > 1:
+            prompt = prompt * num_images_per_prompt
+            bboxes = torch.cat(
+                [bboxes] * num_images_per_prompt, dim=0
+            )
+            category_conditions = torch.cat(
+                [category_conditions] * num_images_per_prompt, dim=0
+            )
+            mask_conditions = torch.cat(
+                [mask_conditions] * num_images_per_prompt, dim=0
+            )
+            mask_vector = torch.cat(
+                [mask_vector] * num_images_per_prompt, dim=0
+            )
+
+        total_batch = batch_size * num_images_per_prompt
+
+        # 1. Encode text prompts
+        text_embeddings = self._encode_prompt(prompt)
+
+        # 2. Encode unconditional prompt for CFG
+        if guidance_scale > 1.0:
+            uncond_embeddings = self._encode_prompt([""] * total_batch)
+
+        # 3. Move conditions to device
+        bboxes = bboxes.to(device).float()
+        category_conditions = category_conditions.to(device).float()
+        mask_conditions = mask_conditions.to(device).float()
+        mask_vector = mask_vector.to(device).float()
+
+        # 4. Encode bbox conditions
+        control = self.condition_encoder(
+            text_embeddings=[category_conditions],
+            masks=[mask_vector],
+            boxes=[bboxes],
+        )
+
+        # 5. Prepare latent noise
+        latent_shape = (
+            total_batch,
+            4,
+            height // 8,
+            width // 8,
+        )
+        latents = torch.randn(
+            latent_shape, device=device, generator=generator
+        )
+
+        # 6. Set up scheduler timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        # 7. Scale initial noise by scheduler init_noise_sigma
+        latents = latents * self.scheduler.init_noise_sigma
+
+        # 8. Denoising loop
+        for t in self.scheduler.timesteps:
+            timesteps = torch.full(
+                (total_batch,), t, device=device, dtype=torch.long
+            )
+
+            if guidance_scale > 1.0:
+                # Classifier-free guidance: run model twice
+                latent_input = torch.cat([latents, latents], dim=0)
+                timestep_input = torch.cat([timesteps, timesteps], dim=0)
+
+                context_in = torch.cat(
+                    [uncond_embeddings, text_embeddings], dim=0
+                )
+                control_in = torch.cat([control, control], dim=0)
+                category_in = [
+                    torch.cat(
+                        [category_conditions, category_conditions], dim=0
+                    )
+                ]
+                mask_in = [
+                    torch.cat(
+                        [mask_conditions, mask_conditions], dim=0
+                    )
+                ]
+
+                noise_pred = self.unet(
+                    x=latent_input,
+                    timesteps=timestep_input,
+                    context=context_in,
+                    control=control_in,
+                    category_control=category_in,
+                    mask_control=mask_in,
+                )
+
+                noise_uncond, noise_text = noise_pred.chunk(2)
+                noise_pred = noise_uncond + guidance_scale * (
+                    noise_text - noise_uncond
+                )
+            else:
+                noise_pred = self.unet(
+                    x=latents,
+                    timesteps=timesteps,
+                    context=text_embeddings,
+                    control=control,
+                    category_control=[category_conditions],
+                    mask_control=[mask_conditions],
+                )
+
+            # Use diffusers scheduler step
+            scheduler_output = self.scheduler.step(
+                model_output=noise_pred,
+                timestep=t,
+                sample=latents,
+                eta=eta,
+                generator=generator,
+            )
+            latents = scheduler_output.prev_sample
+
+        # 9. Decode latents
+        images = self._decode_latents(latents)
+
+        # 10. Post-process
+        if output_type == "pil":
+            images = (
+                einops.rearrange(images, "b c h w -> b h w c") * 127.5 + 127.5
+            )
+            images = images.cpu().numpy().clip(0, 255).astype(np.uint8)
+            images = [Image.fromarray(img) for img in images]
+        elif output_type == "tensor":
+            images = images.cpu()
+        else:
+            raise ValueError(
+                f"Unknown output_type '{output_type}'. "
+                "Use 'pil' or 'tensor'."
+            )
+
+        return AeroGenPipelineOutput(images=images)

scheduler/scheduler_config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "DDIMScheduler",
+  "_diffusers_version": "0.37.0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "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,5 @@
+"""AeroGen text encoder (CLIP)."""
+
+from .clip_text_encoder import AeroGenCLIPTextEncoder
+
+__all__ = ["AeroGenCLIPTextEncoder"]

text_encoder/clip_text_encoder.py

@@ -0,0 +1,43 @@
+"""CLIP text encoder for AeroGen. Uses transformers only (no ldm)."""
+
+import torch.nn as nn
+from transformers import CLIPTokenizer, CLIPTextModel
+
+
+class AeroGenCLIPTextEncoder(nn.Module):
+    """CLIP text encoder compatible with FrozenCLIPEmbedder interface.
+    Uses transformers CLIPTextModel + CLIPTokenizer. No ldm dependency.
+    """
+
+    def __init__(self, version: str = "openai/clip-vit-large-patch14", device: str = "cuda", max_length: int = 77):
+        super().__init__()
+        self.tokenizer = CLIPTokenizer.from_pretrained(version)
+        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.device = device
+        self.max_length = max_length
+        self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, text):
+        if isinstance(text, str):
+            text = [text]
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=True,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+        device = next(self.parameters()).device
+        tokens = batch_encoding["input_ids"].to(device)
+        outputs = self.transformer(input_ids=tokens)
+        return outputs.last_hidden_state
+
+    def encode(self, text):
+        return self(text)

text_encoder/config.json

@@ -0,0 +1,6 @@
+{
+  "target": "text_encoder.clip_text_encoder.AeroGenCLIPTextEncoder",
+  "params": {
+    "version": "openai/clip-vit-large-patch14"
+  }
+}

text_encoder/diffusion_pytorch_model.safetensors

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

unet/__init__.py

@@ -0,0 +1,5 @@
+"""AeroGen UNet components."""
+
+from .unet_aerogen import AeroGenUNet2DConditionModel
+
+__all__ = ["AeroGenUNet2DConditionModel"]

unet/attention_dual.py

@@ -0,0 +1,136 @@
+"""SpatialTransformer with control support - self-contained, no ldm."""
+
+from inspect import isfunction
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from einops import rearrange, repeat
+
+from .diffusion_util import checkpoint
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+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.):
+        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 zero_module(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=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_k_control = nn.Linear(context_dim, inner_dim, bias=False)
+        self.to_v_control = 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, control=None, mask=None, lambda_=1):
+        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))
+        sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
+        if exists(mask):
+            k_control = self.to_k_control(control)
+            v_control = self.to_v_control(control)
+            k_control, v_control = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (k_control, v_control))
+            sim_control = einsum("b i d, b j d -> b i j", q, k_control) * self.scale
+            attn_control = sim_control.softmax(dim=-1)
+            out_control = einsum("b i j, b j d -> b i d", attn_control, v_control)
+            out_control = rearrange(out_control, "(b h) n d -> b n (h d)", h=h)
+        attn = sim.softmax(dim=-1)
+        out = einsum("b i j, b j d -> b i d", attn, v)
+        out = rearrange(out, "(b h) n d -> b n (h d)", h=h)
+        if exists(mask):
+            out = out + lambda_ * out_control
+        return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint_use=True):
+        super().__init__()
+        self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout)
+        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+        self.attn2 = CrossAttention(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_use = checkpoint_use
+
+    def forward(self, x, context=None, control=None, mask=None):
+        return checkpoint(self._forward, (x, context, control, mask), self.parameters(), self.checkpoint_use)
+
+    def _forward(self, x, context=None, control=None, mask=None):
+        x = self.attn1(self.norm1(x)) + x
+        x = self.attn2(self.norm2(x), context=context, control=control, mask=mask) + 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., context_dim=None):
+        super().__init__()
+        self.in_channels = in_channels
+        inner_dim = n_heads * d_head
+        self.norm = Normalize(in_channels)
+        self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        self.transformer_blocks = nn.ModuleList(
+            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim) for _ in range(depth)]
+        )
+        self.proj_out = zero_module(nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0))
+
+    def forward(self, x, context=None, control=None, mask=None):
+        b, c, h, w = x.shape
+        x_in = x
+        x = self.norm(x)
+        x = self.proj_in(x)
+        x = rearrange(x, "b c h w -> b (h w) c")
+        for block in self.transformer_blocks:
+            x = block(x, context=context, control=control, mask=mask)
+        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
+        x = self.proj_out(x)
+        return x + x_in

unet/config.json

@@ -0,0 +1,27 @@
+{
+  "target": "unet.unet_aerogen.AeroGenUNet2DConditionModel",
+  "params": {
+    "image_size": 32,
+    "in_channels": 4,
+    "out_channels": 4,
+    "model_channels": 320,
+    "attention_resolutions": [
+      4,
+      2,
+      1
+    ],
+    "num_res_blocks": 2,
+    "channel_mult": [
+      1,
+      2,
+      4,
+      4
+    ],
+    "num_heads": 8,
+    "use_spatial_transformer": true,
+    "transformer_depth": 1,
+    "context_dim": 768,
+    "use_checkpoint": true,
+    "legacy": false
+  }
+}

unet/diffusion_pytorch_model.safetensors

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

unet/diffusion_util.py

@@ -0,0 +1,92 @@
+"""Diffusion util functions - self-contained, no ldm dependency."""
+
+import math
+import torch
+import torch.nn as nn
+from einops import repeat
+
+
+def checkpoint(func, inputs, params, flag):
+    if flag:
+        args = tuple(inputs) + tuple(params)
+        return _CheckpointFunction.apply(func, len(inputs), *args)
+    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:])
+        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():
+            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 timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+    if not repeat_only:
+        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)
+    else:
+        embedding = repeat(timesteps, "b -> b d", d=dim)
+    return embedding
+
+
+def zero_module(module):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def normalization(channels):
+    return GroupNorm32(32, channels)
+
+
+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 avg_pool_nd(dims, *args, **kwargs):
+    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}")

unet/mask_attention.py

@@ -0,0 +1,113 @@
+"""MaskCrossAttention - self-contained, no ldm/bldm."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import einsum
+from einops import rearrange, repeat
+
+from .diffusion_util import zero_module
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if val is not None else d
+
+
+def Normalize(in_channels):
+    return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+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.):
+        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)
+
+
+class CrossAttention(nn.Module):
+    """Simple cross-attention for MaskCrossAttention (no control branch)."""
+
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        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)
+        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))
+        sim = einsum("b i d, b j d -> b i j", q, k) * self.scale
+        if exists(mask):
+            mask = rearrange(mask, "b ... -> b (...)")
+            max_neg_value = -torch.finfo(sim.dtype).max
+            mask = repeat(mask, "b j -> (b h) j (c)", h=h, c=sim.shape[-1])
+            sim_copy = sim.clone()
+            sim_copy.masked_fill_(~mask, max_neg_value)
+            sim = sim_copy
+        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 MaskCrossAttention(nn.Module):
+    def __init__(self, in_channels, n_heads, d_head, inner_dim=320, context_dim=None):
+        super().__init__()
+        self.in_channels = in_channels
+        self.inner_dim = inner_dim
+        self.norm1 = Normalize(in_channels)
+        self.norm2 = nn.LayerNorm(inner_dim)
+        self.norm3 = nn.LayerNorm(inner_dim)
+        self.ffn = FeedForward(dim=inner_dim, dim_out=inner_dim, dropout=0.0, glu=True)
+        self.proj_in = nn.Linear(in_channels, 320)
+        self.crossattn = CrossAttention(query_dim=inner_dim, heads=n_heads, dim_head=d_head, context_dim=context_dim)
+        self.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
+
+    def forward(self, x, category_control, mask_control, timesteps, attention_strength=0.2, ts_m=200):
+        ts = timesteps[0].item()
+        if ts < ts_m:
+            return x
+        x_in = x
+        mask_control = mask_control[0]
+        category_control = category_control[0]
+        b, c, h, w = x.shape
+        _, n, _, _ = mask_control.shape
+        x = repeat(x, "b c h w -> (b n) c h w", n=n)
+        mask_control_in = rearrange(mask_control, "b n h w -> (b n) h w").contiguous().bool()
+        category_control = rearrange(category_control.unsqueeze(2), "b n c l -> (b n) c l").contiguous()
+        x = rearrange(x, "(b n) c h w -> (b n) (h w) c", b=b, n=n).contiguous()
+        x = self.proj_in(x)
+        x = self.crossattn(self.norm2(x), category_control, mask_control_in) + x
+        x = self.ffn(self.norm3(x)) + x
+        x = self.proj_out(x)
+        x = rearrange(x, "(b n) (h w) c -> b n c h w", b=b, n=n, h=h, w=w).contiguous()
+        mask_control = mask_control.unsqueeze(2).expand(-1, -1, c, -1, -1)
+        x = x * mask_control
+        x_sum = x.sum(dim=1)
+        return attention_strength * x_sum + (1 - attention_strength) * x_in

unet/openaimodel_bbox.py

@@ -0,0 +1,761 @@
+# attn+maskattn+noise
+
+from abc import abstractmethod
+from functools import partial
+import math
+from typing import Iterable
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .diffusion_util import (
+    checkpoint,
+    conv_nd,
+    linear,
+    avg_pool_nd,
+    zero_module,
+    normalization,
+    timestep_embedding,
+)
+from .attention_dual import SpatialTransformer
+from .mask_attention import MaskCrossAttention
+
+
+# dummy replace
+def convert_module_to_f16(x):
+    pass
+
+def convert_module_to_f32(x):
+    pass
+
+
+## go
+class AttentionPool2d(nn.Module):
+    """
+    Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
+    """
+
+    def __init__(
+        self,
+        spacial_dim: int,
+        embed_dim: int,
+        num_heads_channels: int,
+        output_dim: int = None,
+    ):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
+        self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
+        self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
+        self.num_heads = embed_dim // num_heads_channels
+        self.attention = QKVAttention(self.num_heads)
+
+    def forward(self, x):
+        b, c, *_spatial = x.shape
+        x = x.reshape(b, c, -1)  # NC(HW)
+        x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)  # NC(HW+1)
+        x = x + self.positional_embedding[None, :, :].to(x.dtype)  # NC(HW+1)
+        x = self.qkv_proj(x)
+        x = self.attention(x)
+        x = self.c_proj(x)
+        return x[:, :, 0]
+
+
+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):
+    """
+    A sequential module that passes timestep embeddings to the children that
+    support it as an extra input.
+    """
+
+    def forward(self, x, emb, context=None, control=None, mask=None):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            elif isinstance(layer, SpatialTransformer):
+                x = layer(x, context, control = control, mask = mask)
+            else:
+                x = layer(x)
+        return x
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 upsampling occurs in the inner-two dimensions.
+    """
+
+    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 TransposedUpsample(nn.Module):
+    'Learned 2x upsampling without padding'
+    def __init__(self, channels, out_channels=None, ks=5):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+
+        self.up = nn.ConvTranspose2d(self.channels,self.out_channels,kernel_size=ks,stride=2)
+
+    def forward(self,x):
+        return self.up(x)
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+                 downsampling occurs in the inner-two dimensions.
+    """
+
+    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):
+    """
+    A residual block that can optionally change the number of channels.
+    :param channels: the number of input channels.
+    :param emb_channels: the number of timestep embedding channels.
+    :param dropout: the rate of dropout.
+    :param out_channels: if specified, the number of out channels.
+    :param use_conv: if True and out_channels is specified, use a spatial
+        convolution instead of a smaller 1x1 convolution to change the
+        channels in the skip connection.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param use_checkpoint: if True, use gradient checkpointing on this module.
+    :param up: if True, use this block for upsampling.
+    :param down: if True, use this block for downsampling.
+    """
+
+    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):
+        """
+        Apply the block to a Tensor, conditioned on a timestep embedding.
+        :param x: an [N x C x ...] Tensor of features.
+        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        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 = th.chunk(emb_out, 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):
+    """
+    An attention block that allows spatial positions to attend to each other.
+    Originally ported from here, but adapted to the N-d case.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+    """
+
+    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
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            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)
+        if use_new_attention_order:
+            # split qkv before split heads
+            self.attention = QKVAttention(self.num_heads)
+        else:
+            # split heads before split qkv
+            self.attention = 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)   # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
+        #return pt_checkpoint(self._forward, x)  # pytorch
+
+    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)
+
+
+def count_flops_attn(model, _x, y):
+    """
+    A counter for the `thop` package to count the operations in an
+    attention operation.
+    Meant to be used like:
+        macs, params = thop.profile(
+            model,
+            inputs=(inputs, timestamps),
+            custom_ops={QKVAttention: QKVAttention.count_flops},
+        )
+    """
+    b, c, *spatial = y[0].shape
+    num_spatial = int(np.prod(spatial))
+    # We perform two matmuls with the same number of ops.
+    # The first computes the weight matrix, the second computes
+    # the combination of the value vectors.
+    matmul_ops = 2 * b * (num_spatial ** 2) * c
+    model.total_ops += th.DoubleTensor([matmul_ops])
+
+
+class QKVAttentionLegacy(nn.Module):
+    """
+    A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        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 = th.einsum(
+            "bct,bcs->bts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v)
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class QKVAttention(nn.Module):
+    """
+    A module which performs QKV attention and splits in a different order.
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv):
+        """
+        Apply QKV attention.
+        :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        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 = th.einsum(
+            "bct,bcs->bts",
+            (q * scale).view(bs * self.n_heads, ch, length),
+            (k * scale).view(bs * self.n_heads, ch, length),
+        )  # More stable with f16 than dividing afterwards
+        weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+        a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+        return a.reshape(bs, -1, length)
+
+    @staticmethod
+    def count_flops(model, _x, y):
+        return count_flops_attn(model, _x, y)
+
+
+class UNetModel(nn.Module):
+    """
+    The full UNet model with attention and timestep embedding.
+    :param in_channels: channels in the input Tensor.
+    :param model_channels: base channel count for the model.
+    :param out_channels: channels in the output Tensor.
+    :param num_res_blocks: number of residual blocks per downsample.
+    :param attention_resolutions: a collection of downsample rates at which
+        attention will take place. May be a set, list, or tuple.
+        For example, if this contains 4, then at 4x downsampling, attention
+        will be used.
+    :param dropout: the dropout probability.
+    :param channel_mult: channel multiplier for each level of the UNet.
+    :param conv_resample: if True, use learned convolutions for upsampling and
+        downsampling.
+    :param dims: determines if the signal is 1D, 2D, or 3D.
+    :param num_classes: if specified (as an int), then this model will be
+        class-conditional with `num_classes` classes.
+    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
+    :param num_heads: the number of attention heads in each attention layer.
+    :param num_heads_channels: if specified, ignore num_heads and instead use
+                               a fixed channel width per attention head.
+    :param num_heads_upsample: works with num_heads to set a different number
+                               of heads for upsampling. Deprecated.
+    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+    :param resblock_updown: use residual blocks for up/downsampling.
+    :param use_new_attention_order: use a different attention pattern for potentially
+                                    increased efficiency.
+    """
+
+    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,    # custom transformer support
+        transformer_depth=1,              # custom transformer support
+        context_dim=None,                 # custom transformer support
+        n_embed=None,                     # custom support for prediction of discrete ids into codebook of first stage vq model
+        legacy=True,
+
+
+    ):
+        super().__init__()
+        if use_spatial_transformer:
+            assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
+
+        if context_dim is not None:
+            assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
+            if hasattr(context_dim, '__iter__') and not isinstance(context_dim, (int, float, str)):
+                context_dim = list(context_dim)
+
+        if num_heads_upsample == -1:
+            num_heads_upsample = num_heads
+
+        if num_heads == -1:
+            assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
+
+        if num_head_channels == -1:
+            assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
+
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        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 = th.float16 if use_fp16 else th.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 self.num_classes is not None:
+            self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+
+        self.maskcrossattention=MaskCrossAttention(in_channels = 320,n_heads=8,
+                                         d_head=40,context_dim=768)
+
+        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 _ in range(num_res_blocks):
+                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 = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        #num_heads = 1
+                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+                    layers.append(
+                        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
+                        )
+                    )
+                self.input_blocks.append(TimestepEmbedSequential(*layers))
+                self._feature_size += ch
+                input_block_chans.append(ch)
+            if level != len(channel_mult) - 1:
+                out_ch = ch
+                self.input_blocks.append(
+                    TimestepEmbedSequential(
+                        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
+                        )
+                    )
+                )
+                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 = ch // num_head_channels
+            dim_head = num_head_channels
+        if legacy:
+            #num_heads = 1
+            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+        self.middle_block = TimestepEmbedSequential(
+            ResBlock(
+                ch,
+                time_embed_dim,
+                dropout,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+            ),
+            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
+                        ),
+            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(num_res_blocks + 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 = ch // num_head_channels
+                        dim_head = num_head_channels
+                    if legacy:
+                        #num_heads = 1
+                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+                    layers.append(
+                        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
+                        )
+                    )
+                if level and i == num_res_blocks:
+                    out_ch = ch
+                    layers.append(
+                        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)
+                    )
+                    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),
+            #nn.LogSoftmax(dim=1)  # change to cross_entropy and produce non-normalized logits
+        )
+
+    def convert_to_fp16(self):
+        """
+        Convert the torso of the model to float16.
+        """
+        self.input_blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+        self.output_blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self):
+        """
+        Convert the torso of the model to float32.
+        """
+        self.input_blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+        self.output_blocks.apply(convert_module_to_f32)
+
+
+    #x=x_noisy, timesteps=t, context=control,mask_control=mask_control,masks=mask_vector
+
+    def forward(self, x, timesteps=None, context=None, control=None, category_control=None, mask_control=None, y=None,**kwargs):
+        """
+        Apply the model to an input batch.
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :param context: conditioning plugged in via crossattn
+        :param y: an [N] Tensor of labels, if class-conditional.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        assert (y is not None) == (
+            self.num_classes is not None
+        ), "must specify y if and only if the model is class-conditional"
+        hs = []
+        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+        emb = self.time_embed(t_emb)
+
+        # add noise mask
+        mask = th.max(mask_control[0], dim=1).values
+
+        if self.num_classes is not None:
+            assert y.shape == (x.shape[0],)
+            emb = emb + self.label_emb(y)
+
+        h = x.type(self.dtype)
+        for module in self.input_blocks:
+            if mask_control is not None:
+                h = module(h, emb, context, control, mask)
+                h=self.maskcrossattention(h,category_control=category_control,mask_control=mask_control,timesteps=timesteps)
+                category_control=None
+                mask_control=None
+            else:
+                h = module(h, emb, context, control, mask)
+            hs.append(h)
+        h = self.middle_block(h, emb, context, control, mask)
+        for module in self.output_blocks:
+            h = th.cat([h, hs.pop()], dim=1)
+            h = module(h, emb, context, control, mask)
+        h = h.type(x.dtype)
+        if self.predict_codebook_ids:
+            return self.id_predictor(h)
+        else:
+            return self.out(h)

unet/unet_aerogen.py

@@ -0,0 +1,73 @@
+"""
+AeroGen UNet: diffusers ModelMixin wrapper for the custom bbox-conditioned UNet.
+
+Self-contained - no ldm/bldm dependency. Uses local openaimodel_bbox.UNetModel.
+"""
+
+from diffusers import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+from .openaimodel_bbox import UNetModel
+
+
+class AeroGenUNet2DConditionModel(ModelMixin, ConfigMixin):
+    """
+    Diffusers-compatible wrapper for AeroGen's bbox-conditioned UNet.
+    Forward signature: x, timesteps, context, control, category_control, mask_control.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        image_size: int = 32,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        model_channels: int = 320,
+        attention_resolutions: tuple = (4, 2, 1),
+        num_res_blocks: int = 2,
+        channel_mult: tuple = (1, 2, 4, 4),
+        num_heads: int = 8,
+        use_spatial_transformer: bool = True,
+        transformer_depth: int = 1,
+        context_dim: int = 768,
+        use_checkpoint: bool = True,
+        legacy: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.model = UNetModel(
+            image_size=image_size,
+            in_channels=in_channels,
+            model_channels=model_channels,
+            out_channels=out_channels,
+            num_res_blocks=num_res_blocks,
+            attention_resolutions=list(attention_resolutions),
+            channel_mult=list(channel_mult),
+            num_heads=num_heads,
+            use_spatial_transformer=use_spatial_transformer,
+            transformer_depth=transformer_depth,
+            context_dim=context_dim,
+            use_checkpoint=use_checkpoint,
+            legacy=legacy,
+            **kwargs,
+        )
+
+    def forward(
+        self,
+        x,
+        timesteps,
+        context=None,
+        control=None,
+        category_control=None,
+        mask_control=None,
+        **kwargs,
+    ):
+        return self.model(
+            x,
+            timesteps,
+            context=context,
+            control=control,
+            category_control=category_control or [],
+            mask_control=mask_control or [],
+            **kwargs,
+        )

vae/config.json

@@ -0,0 +1,27 @@
+{
+  "target": "ldm.models.autoencoder.AutoencoderKL",
+  "params": {
+    "embed_dim": 4,
+    "monitor": "val/rec_loss",
+    "ddconfig": {
+      "double_z": true,
+      "z_channels": 4,
+      "resolution": 256,
+      "in_channels": 3,
+      "out_ch": 3,
+      "ch": 128,
+      "ch_mult": [
+        1,
+        2,
+        4,
+        4
+      ],
+      "num_res_blocks": 2,
+      "attn_resolutions": [],
+      "dropout": 0.0
+    },
+    "lossconfig": {
+      "target": "torch.nn.Identity"
+    }
+  }
+}

vae/diffusion_pytorch_model.safetensors

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