Upload folder using huggingface_hub

README.md

@@ -0,0 +1,267 @@
+# IntrinsicWeather (Diffusers)
+
+Diffusers-format checkpoint for **[IntrinsicWeather: Controllable Weather Editing in Intrinsic Space](https://arxiv.org/pdf/2508.06982v6)** (CVPR 2026 Highlight).
+
+This repo bundles inverse rendering, forward weather rendering, and the IMAA gating module into a single Hugging Face–compatible layout. Shared Stable Diffusion 3 components (VAE, text encoders, tokenizers, scheduler) are stored once; task-specific transformers live under `transformer/<variant>/`.
+
+## Model layout
+
+```
+IntrisicWeather-diffusers/
+├── dinov2/                         # bundled DINOv2 weights (for IMAA / decomposition)
+├── imaa/                           # Intrinsic Map-Aware Attention weights
+├── text_encoder/, text_encoder_2/, text_encoder_3/
+├── tokenizer/, tokenizer_2/, tokenizer_3/
+├── vae/, scheduler/
+├── transformer/
+│   ├── inverse-512/                # IntrinsicWeatherSD3Transformer2DModel (in_channels=32)
+│   │   └── transformer_intrinsic_weather.py
+│   └── forward/                    # SD3Transformer2DModel (in_channels=96)
+│       └── lora/                   # forward-renderer LoRA (loaded by default)
+├── pipeline_intrinsic_weather.py           # unified: RGB → maps → weather RGB
+├── pipeline_intrinsic_weather_inverse.py   # inverse only
+├── pipeline_intrinsic_weather_forward.py   # forward only
+├── pipeline_utils.py
+├── model_index.json
+├── convert_inverse_renderer_512.py
+├── convert_forward_renderer.py
+└── test_all_pipelines.py
+```
+
+| Component | Source | Notes |
+|-----------|--------|-------|
+| Inverse transformer | [GilgameshYX/InverseRenderer-512](https://huggingface.co/GilgameshYX/InverseRenderer-512) | 512×512 decomposition |
+| Forward transformer + LoRA | [GilgameshYX/ForwardRenderer](https://huggingface.co/GilgameshYX/ForwardRenderer) | LoRA in `transformer/forward/lora/` |
+| IMAA | InverseRenderer-512 `imaa.pth` | Required for map-aware inverse attention |
+| SD3 shared weights | `stabilityai/stable-diffusion-3-medium-diffusers` | VAE + text encoders only |
+| Transformer config | `stabilityai/stable-diffusion-3.5-medium` | Architecture template for weight loading |
+
+## Requirements
+
+- Python 3.10+
+- CUDA GPU recommended (~20 GB VRAM for full end-to-end inference at 512×512)
+- `torch`, `diffusers>=0.38`, `transformers`, `safetensors`, `torchvision`, `Pillow`
+
+```bash
+pip install torch diffusers transformers safetensors torchvision pillow accelerate
+```
+
+## Quick start (end-to-end weather edit)
+
+The unified pipeline decomposes an input RGB image into intrinsic maps, then renders a weather-conditioned result. **DINOv2** is required for decomposition (bundled under `dinov2/`, or use `facebook/dinov2-base` from Hugging Face).
+
+```python
+from pathlib import Path
+
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModel
+
+from pipeline_intrinsic_weather import IntrinsicWeatherPipeline
+
+repo_dir = Path(".").resolve()  # path to this folder
+device = "cuda"
+dtype = torch.bfloat16
+
+pipe = IntrinsicWeatherPipeline.from_pretrained(
+    repo_dir,
+    inverse_transformer_subfolder="inverse-512",
+    forward_transformer_subfolder="forward",
+    device=device,
+    local_files_only=True,
+    torch_dtype=dtype,
+    load_lora=True,
+    load_imaa=True,
+)
+
+dino_path = repo_dir / "dinov2"
+dino_processor = AutoImageProcessor.from_pretrained(dino_path, local_files_only=True)
+dino_model = AutoModel.from_pretrained(dino_path, local_files_only=True).to(device)
+dino_model.eval()
+
+image = Image.open("input.png").convert("RGB")
+result = pipe(
+    image=image,
+    weather="snowy",          # rainy | sunny | snowy | foggy | overcast | night
+    dino_model=dino_model,
+    dino_processor=dino_processor,
+    image_size=512,
+    render_size=512,
+    num_inverse_steps=50,
+    num_forward_steps=50,
+    guidance_scale=6.0,
+    image_guidance_scale=1.5,
+    generator=torch.Generator(device=device).manual_seed(42),
+)
+result.images[0].save("output_snowy.png")
+```
+
+Run from inside this directory (or add it to `PYTHONPATH`) so `pipeline_intrinsic_weather.py` and `imaa/` resolve correctly.
+
+## Pipelines
+
+### 1. `IntrinsicWeatherPipeline` (unified)
+
+Full pipeline: **RGB → intrinsic maps → weather RGB**.
+
+```python
+pipe = IntrinsicWeatherPipeline.from_pretrained(
+    repo_dir,
+    inverse_transformer_subfolder="inverse-512",
+    forward_transformer_subfolder="forward",
+    device="cuda",
+    torch_dtype=torch.bfloat16,
+)
+```
+
+Useful kwargs:
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `inverse_transformer_subfolder` | `"inverse-512"` | Inverse transformer under `transformer/` |
+| `forward_transformer_subfolder` | `"forward"` | Forward transformer under `transformer/` |
+| `load_lora` | `True` | Load LoRA from `transformer/forward/lora/` |
+| `load_imaa` | `True` | Load IMAA weights from `imaa/` |
+| `device` | `None` | Moves all modules to device (IMAA stays float32) |
+
+Sub-methods:
+
+- `pipe.decompose(image, dino_model, dino_processor, ...)` → dict of intrinsic maps
+- `pipe.render(maps, weather="rainy", ...)` → weather-conditioned RGB
+
+### 2. `IntrinsicWeatherInversePipeline`
+
+Inverse rendering only (single intrinsic map per call).
+
+```python
+from pipeline_intrinsic_weather_inverse import IntrinsicWeatherInversePipeline
+
+pipe = IntrinsicWeatherInversePipeline.from_pretrained(
+    repo_dir,
+    transformer_subfolder="inverse-512",
+    device="cuda",
+    torch_dtype=torch.bfloat16,
+)
+```
+
+Load the transformer separately if needed:
+
+```python
+transformer = IntrinsicWeatherInversePipeline.load_transformer(
+    "inverse-512", repo_dir, device="cuda"
+)
+pipe = IntrinsicWeatherInversePipeline.from_pretrained(
+    repo_dir, transformer=transformer, device="cuda"
+)
+```
+
+IMAA and DINO are used by the unified pipeline’s `decompose()` path; for standalone inverse calls, pass `map_aware_mask` from IMAA manually (see `test_all_pipelines.py`).
+
+### 3. `IntrinsicWeatherForwardPipeline`
+
+Forward weather rendering from intrinsic maps.
+
+```python
+from pipeline_intrinsic_weather_forward import IntrinsicWeatherForwardPipeline
+
+pipe = IntrinsicWeatherForwardPipeline.from_pretrained(
+    repo_dir,
+    transformer_subfolder="forward",
+    device="cuda",
+    torch_dtype=torch.bfloat16,
+    load_lora=True,
+)
+```
+
+LoRA weights are read from `transformer/forward/lora/` when `load_lora=True`.
+
+## Weather presets
+
+Built-in weather keys (or pass a custom prompt string):
+
+| Key | Prompt |
+|-----|--------|
+| `rainy` | A rainy day. |
+| `sunny` | A sunny day. |
+| `snowy` | A snowy day. |
+| `foggy` | A foggy day. |
+| `overcast` | An overcast day. |
+| `night` | A night scene. |
+
+## Intrinsic maps (AoVs)
+
+The inverse renderer produces five appearance-of-variety maps:
+
+`albedo`, `normal`, `roughness`, `metallic`, `irradiance`
+
+## Loading transformers manually
+
+Transformers are stored per variant under `transformer/<subfolder>/`. Use `pipeline_utils.load_transformer_from_subfolder`:
+
+```python
+from pipeline_utils import load_transformer_from_subfolder, load_transformer_lora
+
+inverse = load_transformer_from_subfolder(repo_dir, "inverse-512", device="cuda")
+forward = load_transformer_from_subfolder(repo_dir, "forward", device="cuda")
+```
+
+- `inverse-512` uses a custom `IntrinsicWeatherSD3Transformer2DModel` (`in_channels=32`).
+- `forward` uses the standard `SD3Transformer2DModel` (`in_channels=96`).
+
+## Dtype and device notes
+
+- Default dtype is **`torch.bfloat16`** for transformers, VAE, and text encoders.
+- **IMAA** stays in **float32** (DINO patch tokens are float32).
+- Pass `device="cuda"` to `from_pretrained` on all three pipeline classes; the unified pipeline moves every registered module to the target device automatically.
+
+## Testing
+
+Smoke-test all pipelines on CUDA:
+
+```bash
+python test_all_pipelines.py
+```
+
+Runs 2-step inverse, forward (with LoRA), and unified load checks with `bfloat16`.
+
+## Re-converting from original checkpoints
+
+If you have the raw GilgameshYX checkpoints:
+
+```bash
+# Inverse renderer (512) + IMAA
+python convert_inverse_renderer_512.py
+
+# Forward renderer + LoRA
+python convert_forward_renderer.py
+```
+
+See `conversion_metadata.json` and `conversion_metadata_forward.json` for source paths used during conversion.
+
+## Hugging Face Hub loading
+
+When published to the Hub, load with `trust_remote_code=True`:
+
+```python
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained(
+    "BiliSakura/IntrisicWeather-diffusers",
+    custom_pipeline="pipeline_intrinsic_weather.py",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+)
+```
+
+For local use, importing `IntrinsicWeatherPipeline` directly (as in Quick start) is simpler and avoids Hub cache path issues with custom modules.
+
+## References
+
+- **Paper:** [IntrinsicWeather (arXiv:2508.06982)](https://arxiv.org/pdf/2508.06982v6)
+- **Project page:** https://yixinzhu042.github.io/IntrinsicWeather/
+- **Upstream diffusers repo:** [IntrinsicWeather-diffusers](https://github.com/YixinZhu042/IntrinsicWeather)
+- **Original weights:** [GilgameshYX/InverseRenderer-512](https://huggingface.co/GilgameshYX/InverseRenderer-512), [GilgameshYX/ForwardRenderer](https://huggingface.co/GilgameshYX/ForwardRenderer)
+
+## License
+
+Weights and code follow the licenses of the upstream IntrinsicWeather project and the Stable Diffusion 3 components used for shared modules.

__init__.py

@@ -0,0 +1,9 @@
+from pipeline_intrinsic_weather import IntrinsicWeatherPipeline
+from pipeline_intrinsic_weather_forward import IntrinsicWeatherForwardPipeline
+from pipeline_intrinsic_weather_inverse import IntrinsicWeatherInversePipeline
+
+__all__ = [
+    "IntrinsicWeatherPipeline",
+    "IntrinsicWeatherForwardPipeline",
+    "IntrinsicWeatherInversePipeline",
+]

convert_forward_renderer.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+"""Convert GilgameshYX ForwardRenderer into BiliSakura IntrisicWeather-diffusers layout."""
+
+from __future__ import annotations
+
+import json
+import shutil
+import sys
+from pathlib import Path
+
+from diffusers.models.transformers import SD3Transformer2DModel
+
+COLLECTION_ROOT = Path(__file__).resolve().parent
+INTRINSIC_REPO = Path("/data/projects/IntrinsicWeather-diffusers")
+sys.path.insert(0, str(INTRINSIC_REPO / "src"))
+sys.path.insert(0, str(INTRINSIC_REPO))
+
+from scripts._conversion_utils import (  # noqa: E402
+    expand_sd3_input_projection,
+    load_torch,
+    write_scheduler_config,
+)
+
+from _collection_setup import install_hub_pipelines  # noqa: E402
+
+SD3_PATH = Path(
+    "/data/projects/Visual-Generative-Foundation-Model-Collection/models/stabilityai/stable-diffusion-3-medium-diffusers"
+)
+SD35_TRANSFORMER_REPO = "stabilityai/stable-diffusion-3.5-medium"
+CKPT_PATH = Path(
+    "/data/projects/Visual-Generative-Foundation-Model-Collection/models/GilgameshYX/ForwardRenderer"
+)
+OUTPUT_ROOT = COLLECTION_ROOT
+TRANSFORMER_VARIANT = "forward"
+SHARED_COMPONENTS = (
+    "text_encoder",
+    "text_encoder_2",
+    "text_encoder_3",
+    "tokenizer",
+    "tokenizer_2",
+    "tokenizer_3",
+    "vae",
+    "scheduler",
+)
+
+
+def copy_sd3_shared_components(sd3_path: Path, output_path: Path) -> None:
+    for name in SHARED_COMPONENTS:
+        src = sd3_path / name
+        dst = output_path / name
+        if dst.exists():
+            print(f"Skipping existing shared component: {dst}")
+            continue
+        print(f"Copying {name} ...")
+        shutil.copytree(src, dst)
+
+
+def main() -> None:
+    transformer_dir = OUTPUT_ROOT / "transformer" / TRANSFORMER_VARIANT
+    transformer_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Ensuring shared SD3 components from {SD3_PATH} ...")
+    copy_sd3_shared_components(SD3_PATH, OUTPUT_ROOT)
+    write_scheduler_config(OUTPUT_ROOT)
+    install_hub_pipelines(OUTPUT_ROOT)
+
+    print("Converting forward renderer transformer ...")
+    transformer = SD3Transformer2DModel.from_config(
+        SD3Transformer2DModel.load_config(SD35_TRANSFORMER_REPO, subfolder="transformer")
+    )
+    transformer = expand_sd3_input_projection(transformer, in_channels=96)
+    transformer.load_state_dict(load_torch(CKPT_PATH / "pytorch_model.bin"), strict=True)
+    transformer.save_pretrained(transformer_dir.as_posix(), safe_serialization=True)
+
+    print("Saving LoRA weights ...")
+    lora_dir = transformer_dir / "lora"
+    lora_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(CKPT_PATH / "pytorch_lora_weights.safetensors", lora_dir / "pytorch_lora_weights.safetensors")
+
+    conversion_metadata = {
+        "task": "forward_renderer",
+        "transformer_variant": TRANSFORMER_VARIANT,
+        "source_transformer_checkpoint": str((CKPT_PATH / "pytorch_model.bin").resolve()),
+        "source_lora_checkpoint": str((CKPT_PATH / "pytorch_lora_weights.safetensors").resolve()),
+        "lora_dir": str((lora_dir).resolve()),
+        "sd3_path": str(SD3_PATH.resolve()),
+        "sd35_transformer_repo": SD35_TRANSFORMER_REPO,
+        "in_channels": 96,
+    }
+    (OUTPUT_ROOT / "conversion_metadata_forward.json").write_text(
+        json.dumps(conversion_metadata, indent=2) + "\n",
+        encoding="utf-8",
+    )
+    print(f"Saved transformer to: {transformer_dir}")
+    print("Load with: load_forward_pipeline(transformer_subfolder='forward')")
+
+
+if __name__ == "__main__":
+    main()

convert_inverse_renderer_512.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""Convert GilgameshYX InverseRenderer-512 into BiliSakura IntrisicWeather-diffusers layout."""
+
+from __future__ import annotations
+
+import json
+import shutil
+import sys
+from pathlib import Path
+
+from diffusers.models.transformers import SD3Transformer2DModel
+
+COLLECTION_ROOT = Path(__file__).resolve().parent
+INTRINSIC_REPO = Path("/data/projects/IntrinsicWeather-diffusers")
+sys.path.insert(0, str(INTRINSIC_REPO / "src"))
+sys.path.insert(0, str(INTRINSIC_REPO))
+
+from intrinsic_weather.models.transformers.transformer_intrinsic_weather import (  # noqa: E402
+    IntrinsicWeatherSD3Transformer2DModel,
+)
+from scripts._conversion_utils import (  # noqa: E402
+    ROOT as REPO_ROOT,
+    expand_sd3_input_projection,
+    merge_sharded_state_dict,
+    save_imaa_bundle,
+    write_scheduler_config,
+)
+
+from _collection_setup import install_hub_pipelines  # noqa: E402
+
+SD3_PATH = Path(
+    "/data/projects/Visual-Generative-Foundation-Model-Collection/models/stabilityai/stable-diffusion-3-medium-diffusers"
+)
+SD35_TRANSFORMER_REPO = "stabilityai/stable-diffusion-3.5-medium"
+CKPT_PATH = Path(
+    "/data/projects/Visual-Generative-Foundation-Model-Collection/models/GilgameshYX/InverseRenderer-512"
+)
+OUTPUT_ROOT = COLLECTION_ROOT
+TRANSFORMER_VARIANT = "inverse-512"
+SHARED_COMPONENTS = (
+    "text_encoder",
+    "text_encoder_2",
+    "text_encoder_3",
+    "tokenizer",
+    "tokenizer_2",
+    "tokenizer_3",
+    "vae",
+    "scheduler",
+)
+
+
+def copy_sd3_shared_components(sd3_path: Path, output_path: Path) -> None:
+    for name in SHARED_COMPONENTS:
+        src = sd3_path / name
+        dst = output_path / name
+        if dst.exists():
+            print(f"Skipping existing shared component: {dst}")
+            continue
+        print(f"Copying {name} ...")
+        shutil.copytree(src, dst)
+
+
+def main() -> None:
+    transformer_dir = OUTPUT_ROOT / "transformer" / TRANSFORMER_VARIANT
+    transformer_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Copying shared SD3 components from {SD3_PATH} ...")
+    copy_sd3_shared_components(SD3_PATH, OUTPUT_ROOT)
+    write_scheduler_config(OUTPUT_ROOT)
+    install_hub_pipelines(OUTPUT_ROOT)
+
+    print("Converting inverse renderer transformer ...")
+    base_transformer = SD3Transformer2DModel.from_config(
+        SD3Transformer2DModel.load_config(SD35_TRANSFORMER_REPO, subfolder="transformer")
+    )
+    base_transformer = expand_sd3_input_projection(base_transformer, in_channels=32)
+    custom_blocks = IntrinsicWeatherSD3Transformer2DModel.from_config(base_transformer.config)
+    custom_blocks.load_state_dict(
+        merge_sharded_state_dict(
+            [
+                CKPT_PATH / "pytorch_model-00001-of-00002.bin",
+                CKPT_PATH / "pytorch_model-00002-of-00002.bin",
+            ]
+        ),
+        strict=True,
+    )
+    custom_blocks.save_pretrained(transformer_dir.as_posix(), safe_serialization=True)
+    shutil.copy2(
+        REPO_ROOT / "src" / "intrinsic_weather" / "models" / "transformers" / "transformer_intrinsic_weather.py",
+        transformer_dir / "transformer_intrinsic_weather.py",
+    )
+
+    print("Saving IMAA weights ...")
+    save_imaa_bundle(CKPT_PATH / "imaa.pth", OUTPUT_ROOT, safe_serialization=True)
+
+    conversion_metadata = {
+        "task": "inverse_renderer",
+        "resolution": 512,
+        "transformer_variant": TRANSFORMER_VARIANT,
+        "source_transformer_checkpoints": [
+            str((CKPT_PATH / "pytorch_model-00001-of-00002.bin").resolve()),
+            str((CKPT_PATH / "pytorch_model-00002-of-00002.bin").resolve()),
+        ],
+        "source_imaa_checkpoint": str((CKPT_PATH / "imaa.pth").resolve()),
+        "sd3_path": str(SD3_PATH.resolve()),
+        "sd35_transformer_repo": SD35_TRANSFORMER_REPO,
+        "in_channels": 32,
+    }
+    (OUTPUT_ROOT / "conversion_metadata.json").write_text(
+        json.dumps(conversion_metadata, indent=2) + "\n",
+        encoding="utf-8",
+    )
+    print(f"Saved transformer to: {transformer_dir}")
+    print("Load with: load_inverse_pipeline(transformer_subfolder='inverse-512')")
+
+
+if __name__ == "__main__":
+    main()

dinov2/README.md

@@ -0,0 +1,61 @@
+---
+license: apache-2.0
+tags:
+- dino
+- vision
+inference: false
+---
+
+# Vision Transformer (base-sized model) trained using DINOv2
+
+Vision Transformer (ViT) model trained using the DINOv2 method. It was introduced in the paper [DINOv2: Learning Robust Visual Features without Supervision](https://arxiv.org/abs/2304.07193) by Oquab et al. and first released in [this repository](https://github.com/facebookresearch/dinov2). 
+
+Disclaimer: The team releasing DINOv2 did not write a model card for this model so this model card has been written by the Hugging Face team.
+
+## Model description
+
+The Vision Transformer (ViT) is a transformer encoder model (BERT-like) pretrained on a large collection of images in a self-supervised fashion. 
+
+Images are presented to the model as a sequence of fixed-size patches, which are linearly embedded. One also adds a [CLS] token to the beginning of a sequence to use it for classification tasks. One also adds absolute position embeddings before feeding the sequence to the layers of the Transformer encoder.
+
+Note that this model does not include any fine-tuned heads. 
+
+By pre-training the model, it learns an inner representation of images that can then be used to extract features useful for downstream tasks: if you have a dataset of labeled images for instance, you can train a standard classifier by placing a linear layer on top of the pre-trained encoder. One typically places a linear layer on top of the [CLS] token, as the last hidden state of this token can be seen as a representation of an entire image.
+
+## Intended uses & limitations
+
+You can use the raw model for feature extraction. See the [model hub](https://huggingface.co/models?search=facebook/dinov2) to look for
+fine-tuned versions on a task that interests you.
+
+### How to use
+
+Here is how to use this model:
+
+```python
+from transformers import AutoImageProcessor, AutoModel
+from PIL import Image
+import requests
+
+url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+image = Image.open(requests.get(url, stream=True).raw)
+
+processor = AutoImageProcessor.from_pretrained('facebook/dinov2-base')
+model = AutoModel.from_pretrained('facebook/dinov2-base')
+
+inputs = processor(images=image, return_tensors="pt")
+outputs = model(**inputs)
+last_hidden_states = outputs.last_hidden_state
+```
+
+### BibTeX entry and citation info
+
+```bibtex
+misc{oquab2023dinov2,
+      title={DINOv2: Learning Robust Visual Features without Supervision}, 
+      author={Maxime Oquab and Timothée Darcet and Théo Moutakanni and Huy Vo and Marc Szafraniec and Vasil Khalidov and Pierre Fernandez and Daniel Haziza and Francisco Massa and Alaaeldin El-Nouby and Mahmoud Assran and Nicolas Ballas and Wojciech Galuba and Russell Howes and Po-Yao Huang and Shang-Wen Li and Ishan Misra and Michael Rabbat and Vasu Sharma and Gabriel Synnaeve and Hu Xu and Hervé Jegou and Julien Mairal and Patrick Labatut and Armand Joulin and Piotr Bojanowski},
+      year={2023},
+      eprint={2304.07193},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

dinov2/config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "Dinov2Model"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "drop_path_rate": 0.0,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 768,
+  "image_size": 518,
+  "initializer_range": 0.02,
+  "layer_norm_eps": 1e-06,
+  "layerscale_value": 1.0,
+  "mlp_ratio": 4,
+  "model_type": "dinov2",
+  "num_attention_heads": 12,
+  "num_channels": 3,
+  "num_hidden_layers": 12,
+  "patch_size": 14,
+  "qkv_bias": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.31.0.dev0",
+  "use_swiglu_ffn": false
+}

dinov2/model.safetensors

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

dinov2/preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "BitImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 256
+  }
+}

imaa/config.json

@@ -0,0 +1,11 @@
+{
+  "_class_name": "IMAA",
+  "num_maps": 5,
+  "map_embedding_dim": 256,
+  "common_dim": 128,
+  "conv_channels": [
+    128,
+    64
+  ],
+  "dino_patch_dim": 768
+}

imaa/imaa.py

@@ -0,0 +1,205 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def extract_patch_tokens_min_windows(
+    images: torch.Tensor,
+    model: nn.Module,
+    processor,
+    window_size: int = 224,
+    device: str | torch.device = "cuda",
+) -> torch.Tensor:
+    r"""
+    Tile each image with a minimal window set and return averaged DINO patch tokens.
+
+    Args:
+        images (`torch.Tensor`): Batch of RGB images `(B, C, H, W)`.
+        model: DINO vision transformer.
+        processor: Hugging Face image processor for DINO.
+        window_size (`int`): Sliding-window size in pixels.
+        device: Device for intermediate tensors.
+
+    Returns:
+        `torch.Tensor` of shape `(B, H//patch, W//patch, hidden_size)`.
+    """
+    batch_size, _, height, width = images.shape
+    hidden_size = model.config.hidden_size
+    patch_size = model.config.patch_size
+    token_avgs = []
+
+    for batch_idx in range(batch_size):
+        image = images[batch_idx].float()
+        if image.max() <= 1.0:
+            image_np = (image.permute(1, 2, 0).cpu().numpy() * 255).clip(0, 255).astype("uint8")
+        else:
+            image_np = image.permute(1, 2, 0).cpu().numpy().clip(0, 255).astype("uint8")
+
+        token_sum = torch.zeros((height // patch_size, width // patch_size, hidden_size), device=device)
+        token_count = torch.zeros((height // patch_size, width // patch_size, 1), device=device)
+
+        num_y = (height + window_size - 1) // window_size
+        num_x = (width + window_size - 1) // window_size
+        y_positions = [index * window_size for index in range(num_y - 1)] + [height - window_size]
+        x_positions = [index * window_size for index in range(num_x - 1)] + [width - window_size]
+
+        for y in y_positions:
+            for x in x_positions:
+                patch = image_np[y : y + window_size, x : x + window_size, :]
+                inputs = processor(images=patch, return_tensors="pt").to(device)
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                patch_tokens = outputs.last_hidden_state[:, 1:, :]
+                patch_tokens = patch_tokens.reshape(
+                    1, window_size // patch_size, window_size // patch_size, hidden_size
+                ).squeeze(0)
+
+                y0, x0 = y // patch_size, x // patch_size
+                y1, x1 = y0 + window_size // patch_size, x0 + window_size // patch_size
+                token_sum[y0:y1, x0:x1, :] += patch_tokens
+                token_count[y0:y1, x0:x1, 0] += 1
+
+        token_avgs.append(token_sum / token_count)
+
+    return torch.stack(token_avgs, dim=0)
+
+
+class LayerNorm2d(nn.Module):
+    def __init__(self, channels: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm([channels])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        return x.permute(0, 3, 1, 2)
+
+
+class IMAA(nn.Module):
+    r"""
+    Intrinsic Map-Aware Attention (IMAA) gating module.
+
+    Produces per-map attention biases from DINO patch tokens and learnable map embeddings.
+    """
+
+    def __init__(
+        self,
+        dino_model: Optional[nn.Module] = None,
+        processor=None,
+        num_maps: int = 5,
+        map_embedding_dim: int = 256,
+        common_dim: int = 128,
+        conv_channels: Optional[list[int]] = None,
+        dino_patch_dim: int = 768,
+    ) -> None:
+        super().__init__()
+        conv_channels = conv_channels or [128, 64]
+        self.dino = dino_model
+        self.processor = processor
+        if self.dino is not None:
+            self.dino.eval()
+            for param in self.dino.parameters():
+                param.requires_grad = False
+
+        self.num_maps = num_maps
+        self.map_embedding_dim = map_embedding_dim
+        self.common_dim = common_dim
+        self.dino_patch_dim = dino_patch_dim
+        self.map_embedding = nn.Parameter(torch.randn(num_maps, map_embedding_dim))
+        self.dino_proj = nn.Conv2d(dino_patch_dim, common_dim, kernel_size=1)
+        self.map_proj = nn.Linear(map_embedding_dim, common_dim)
+        self.fusion_layer = nn.Sequential(
+            nn.Conv2d(common_dim * 2, common_dim, 1),
+            LayerNorm2d(common_dim),
+            nn.ReLU(),
+            nn.Conv2d(common_dim, common_dim, 3, padding=1),
+        )
+
+        conv_layers: list[nn.Module] = []
+        in_channels = common_dim
+        for out_channels in conv_channels:
+            conv_layers.extend([nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.ReLU()])
+            in_channels = out_channels
+        conv_layers.append(nn.Conv2d(in_channels, 1, kernel_size=1))
+        self.conv_head = nn.Sequential(*conv_layers)
+
+    def forward(
+        self,
+        image: Optional[torch.Tensor] = None,
+        patch_tokens: Optional[torch.Tensor] = None,
+        output_size: Optional[Tuple[int, int]] = None,
+        map_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if patch_tokens is None:
+            if self.dino is None or image is None:
+                raise ValueError("Either `patch_tokens` or (`image` and a frozen DINO model) must be provided.")
+            patch_tokens = extract_patch_tokens_min_windows(
+                image, self.dino, self.processor, window_size=224, device=image.device
+            )
+
+        dino_feat_map = patch_tokens.permute(0, 3, 1, 2)
+        dino_proj = self.dino_proj(dino_feat_map)
+        map_emb = self.map_embedding[map_ids]
+        map_proj = self.map_proj(map_emb).unsqueeze(-1).unsqueeze(-1).expand(-1, -1, dino_proj.size(2), dino_proj.size(3))
+        fused_map = self.fusion_layer(torch.cat([dino_proj, map_proj], dim=1))
+        raw_gating_map = self.conv_head(fused_map)
+        aligned_map = (
+            F.interpolate(raw_gating_map, size=output_size, mode="bilinear", align_corners=False)
+            if output_size is not None
+            else raw_gating_map
+        )
+        return torch.sigmoid(aligned_map)
+
+
+def build_attn_mask(
+    w_gating: torch.Tensor,
+    text_len: int,
+    img_len: int,
+    lam: float,
+) -> torch.Tensor:
+    r"""
+    Build an additive attention mask from IMAA gating weights.
+
+    Args:
+        w_gating (`torch.Tensor`): Gating map `[B, 1, H, W]` or flattened `[B, img_len]`.
+        text_len (`int`): Number of text tokens prepended to image tokens.
+        img_len (`int`): Expected number of image tokens.
+        lam (`float`): Mask scaling factor.
+
+    Returns:
+        Attention bias tensor shaped for SD3 joint attention.
+    """
+    batch_size = w_gating.shape[0]
+    total_len = text_len + img_len
+    if w_gating.dim() == 4:
+        w_gating = w_gating.view(batch_size, -1)
+
+    gating = lam * w_gating
+    actual_img_len = gating.shape[1]
+    if actual_img_len != img_len:
+        if actual_img_len > img_len:
+            gating = gating[:, :img_len]
+        else:
+            padding = torch.zeros(batch_size, img_len - actual_img_len, device=gating.device, dtype=gating.dtype)
+            gating = torch.cat([gating, padding], dim=1)
+
+    col_bias = torch.zeros(batch_size, total_len, device=w_gating.device, dtype=w_gating.dtype)
+    col_bias[:, text_len:] = gating
+    return col_bias.view(batch_size, 1, 1, total_len)

imaa/model.safetensors

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

model_index.json

@@ -0,0 +1,39 @@
+{
+  "_class_name": [
+    "pipeline",
+    "IntrinsicWeatherInversePipeline"
+  ],
+  "_diffusers_version": "0.38.0",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModelWithProjection"
+  ],
+  "text_encoder_2": [
+    "transformers",
+    "CLIPTextModelWithProjection"
+  ],
+  "text_encoder_3": [
+    "transformers",
+    "T5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "tokenizer_2": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "tokenizer_3": [
+    "transformers",
+    "T5TokenizerFast"
+  ]
+}

pipeline_intrinsic_weather.py

@@ -0,0 +1,486 @@
+"""Hub custom pipeline: IntrinsicWeatherPipeline.
+Inverse decomposition + forward weather rendering in one call.
+Load with native Hugging Face diffusers and trust_remote_code=True.
+"""
+
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torchvision.transforms as T
+from pathlib import Path
+from diffusers.image_processor import PipelineImageInput
+from diffusers.loaders import SD3LoraLoaderMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput
+from diffusers.utils import logging, replace_example_docstring
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+from imaa.imaa import IMAA, build_attn_mask, extract_patch_tokens_min_windows
+from pipeline_utils import load_transformer_from_subfolder, load_transformer_lora, resolve_repo_dir
+
+logger = logging.get_logger(__name__)
+
+AOVS = ["albedo", "normal", "roughness", "metallic", "irradiance"]
+INVERSE_PROMPTS = {
+    "albedo": "Albedo (diffuse basecolor)",
+    "normal": "Camera-space Normal",
+    "roughness": "Roughness",
+    "metallic": "Metallicness",
+    "irradiance": "Irradiance (lighting)",
+}
+WEATHER_PROMPTS = {
+    "rainy": "A rainy day.",
+    "sunny": "A sunny day.",
+    "snowy": "A snowy day.",
+    "foggy": "A foggy day.",
+    "overcast": "An overcast day.",
+    "night": "A night scene.",
+}
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from pipeline_intrinsic_weather import IntrinsicWeatherPipeline
+        >>> from transformers import AutoImageProcessor, AutoModel
+
+        >>> repo_dir = Path("./IntrisicWeather-diffusers").resolve()
+        >>> pipe = IntrinsicWeatherPipeline.from_pretrained(
+        ...     repo_dir,
+        ...     inverse_transformer_subfolder="inverse-512",
+        ...     forward_transformer_subfolder="forward",
+        ...     local_files_only=True,
+        ...     torch_dtype=torch.float16,
+        ... )
+        >>> pipe.to("cuda")
+        >>> dino_processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
+        >>> dino_model = AutoModel.from_pretrained("facebook/dinov2-base").to("cuda")
+
+        >>> from PIL import Image
+        >>> image = Image.open("input.png")
+        >>> result = pipe(
+        ...     image=image,
+        ...     weather="rainy",
+        ...     dino_model=dino_model,
+        ...     dino_processor=dino_processor,
+        ... )
+        >>> result.images[0].save("rainy.png")
+        ```
+"""
+
+
+@dataclass
+class IntrinsicWeatherPipelineOutput(StableDiffusion3PipelineOutput):
+    maps: Optional[Dict[str, Union[PIL.Image.Image, np.ndarray]]] = None
+
+
+class IntrinsicWeatherPipeline(DiffusionPipeline, SD3LoraLoaderMixin):
+    r"""
+    End-to-end IntrinsicWeather pipeline: RGB → intrinsic maps → weather-conditioned RGB.
+
+    Parameters:
+        inverse_transformer ([`IntrinsicWeatherSD3Transformer2DModel`]):
+            Map-aware SD3 transformer for intrinsic decomposition.
+        forward_transformer ([`SD3Transformer2DModel`]):
+            SD3 transformer for forward weather rendering.
+        imaa ([`IMAA`]):
+            Map-aware attention gating module used during inverse rendering.
+        vae, text_encoder(s), tokenizer(s), scheduler:
+            Shared Stable Diffusion 3 components.
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->inverse_transformer->forward_transformer->vae"
+    _optional_components: List[str] = []
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        inverse_transformer_subfolder: str = "inverse-512",
+        forward_transformer_subfolder: str = "forward",
+        inverse_transformer=None,
+        forward_transformer=None,
+        load_lora: bool = True,
+        load_imaa: bool = True,
+        **kwargs,
+    ):
+        repo_dir = resolve_repo_dir(pretrained_model_name_or_path)
+        dtype = kwargs.get("torch_dtype", torch.bfloat16)
+        device = kwargs.get("device")
+        if inverse_transformer is None:
+            inverse_transformer = load_transformer_from_subfolder(
+                repo_dir,
+                inverse_transformer_subfolder,
+                dtype=dtype,
+                device=device,
+            )
+        if forward_transformer is None:
+            forward_transformer = load_transformer_from_subfolder(
+                repo_dir,
+                forward_transformer_subfolder,
+                dtype=dtype,
+                device=device,
+            )
+
+        imaa = IMAA(dino_model=None, processor=None, num_maps=5, map_embedding_dim=256, common_dim=128)
+        if device is not None:
+            imaa = imaa.to(device)
+
+        pipe = DiffusionPipeline.from_pretrained(
+            repo_dir.as_posix(),
+            *args,
+            inverse_transformer=inverse_transformer,
+            forward_transformer=forward_transformer,
+            imaa=imaa,
+            custom_pipeline=str(repo_dir / "pipeline_intrinsic_weather.py"),
+            trust_remote_code=True,
+            **kwargs,
+        )
+        if load_imaa:
+            pipe.load_imaa_weights(repo_dir / "imaa")
+        if load_lora:
+            load_transformer_lora(pipe._forward, repo_dir, forward_transformer_subfolder)
+        if device is not None:
+            for name in pipe.components.keys():
+                module = getattr(pipe, name, None)
+                if module is not None and hasattr(module, "to"):
+                    if name == "imaa":
+                        module.to(device=device)
+                    else:
+                        module.to(device=device, dtype=dtype)
+        return pipe
+
+    def __init__(
+        self,
+        inverse_transformer,
+        forward_transformer,
+        imaa: IMAA,
+        scheduler,
+        vae,
+        text_encoder: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer_2: CLIPTokenizer,
+        text_encoder_3: T5EncoderModel,
+        tokenizer_3: T5TokenizerFast,
+    ) -> None:
+        super().__init__()
+        from pipeline_intrinsic_weather_forward import IntrinsicWeatherForwardPipeline
+        from pipeline_intrinsic_weather_inverse import IntrinsicWeatherInversePipeline
+
+        self.register_modules(
+            inverse_transformer=inverse_transformer,
+            forward_transformer=forward_transformer,
+            imaa=imaa,
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            text_encoder_2=text_encoder_2,
+            tokenizer_2=tokenizer_2,
+            text_encoder_3=text_encoder_3,
+            tokenizer_3=tokenizer_3,
+        )
+        shared = dict(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            text_encoder_2=text_encoder_2,
+            tokenizer_2=tokenizer_2,
+            text_encoder_3=text_encoder_3,
+            tokenizer_3=tokenizer_3,
+        )
+        self._inverse = IntrinsicWeatherInversePipeline(transformer=inverse_transformer, **shared)
+        self._forward = IntrinsicWeatherForwardPipeline(transformer=forward_transformer, **shared)
+
+    def load_imaa_weights(self, imaa_dir: Union[str, Path]) -> None:
+        r"""Load IMAA weights from an `imaa/` subfolder produced by the conversion script."""
+        imaa_dir = Path(imaa_dir)
+        weights_path = imaa_dir / "model.safetensors"
+        if weights_path.exists():
+            from safetensors.torch import load_file
+
+            state_dict = load_file(weights_path.as_posix())
+        else:
+            payload = torch.load(imaa_dir / "pytorch_model.bin", map_location="cpu", weights_only=False)
+            state_dict = payload["model_state_dict"] if isinstance(payload, dict) and "model_state_dict" in payload else payload
+        self.imaa.load_state_dict(state_dict)
+        self.imaa.eval()
+
+    @staticmethod
+    def _resolve_weather_prompt(weather: str) -> str:
+        key = weather.lower().strip()
+        if key in WEATHER_PROMPTS:
+            return WEATHER_PROMPTS[key]
+        return weather
+
+    def _preprocess_rgb(self, image: PipelineImageInput, size: int) -> torch.Tensor:
+        if isinstance(image, PIL.Image.Image):
+            pil = image.convert("RGB")
+        elif isinstance(image, torch.Tensor):
+            if image.ndim == 3:
+                image = image.unsqueeze(0)
+            tensor = image
+            if tensor.shape[1] != 3:
+                raise ValueError(f"Expected 3 RGB channels, got shape {tuple(tensor.shape)}")
+            if tensor.min() < 0:
+                tensor = (tensor + 1.0) / 2.0
+            return tensor.to(device=self._execution_device, dtype=self.dtype)
+        else:
+            raise TypeError(f"`image` must be PIL.Image or torch.Tensor, got {type(image)}")
+
+        transform = T.Compose([T.Resize((size, size), interpolation=T.InterpolationMode.BILINEAR), T.ToTensor()])
+        tensor = transform(pil)
+        return tensor.unsqueeze(0).to(device=self._execution_device, dtype=self.dtype)
+
+    @staticmethod
+    def _map_array_to_tensor(
+        array: Union[np.ndarray, PIL.Image.Image, torch.Tensor],
+        size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if isinstance(array, torch.Tensor):
+            tensor = array
+            if tensor.ndim == 3:
+                tensor = tensor.unsqueeze(0)
+            if tensor.min() < 0:
+                tensor = (tensor + 1.0) / 2.0
+            return tensor.to(device=device, dtype=dtype)
+
+        if isinstance(array, np.ndarray):
+            if array.dtype != np.uint8:
+                array = np.clip(array * 255.0, 0, 255).astype(np.uint8)
+            pil = PIL.Image.fromarray(array)
+        else:
+            pil = array.convert("RGB")
+
+        transform = T.Compose(
+            [
+                T.Resize((size, size), interpolation=T.InterpolationMode.BILINEAR),
+                T.ToTensor(),
+            ]
+        )
+        return transform(pil).unsqueeze(0).to(device=device, dtype=dtype)
+
+    @torch.no_grad()
+    def decompose(
+        self,
+        image: PipelineImageInput,
+        dino_model: PreTrainedModel,
+        dino_processor: PreTrainedTokenizer,
+        num_inference_steps: int = 50,
+        image_size: Optional[int] = None,
+        output_type: str = "pil",
+    ) -> Dict[str, Union[PIL.Image.Image, np.ndarray]]:
+        r"""
+        Run inverse rendering and return intrinsic maps.
+
+        Args:
+            image: Input RGB image.
+            dino_model: Frozen DINOv2 vision model.
+            dino_processor: Matching image processor.
+            num_inference_steps: Denoising steps per intrinsic map.
+            image_size: Square resize before decomposition. Defaults to input size.
+            output_type: `"pil"` or `"np"` per map.
+        """
+        if image_size is None:
+            if isinstance(image, PIL.Image.Image):
+                image_size = max(image.size)
+            elif isinstance(image, torch.Tensor):
+                image_size = image.shape[-1]
+            else:
+                image_size = 1024
+
+        image_tensor = self._preprocess_rgb(image, image_size)
+        patch_tokens = extract_patch_tokens_min_windows(
+            image_tensor,
+            dino_model,
+            dino_processor,
+            window_size=224,
+            device=image_tensor.device,
+        )
+        output_size = (image_tensor.shape[2] // 16, image_tensor.shape[3] // 16)
+        img_len = output_size[0] * output_size[1]
+
+        maps: Dict[str, Union[PIL.Image.Image, np.ndarray]] = {}
+        for map_index, aov_name in enumerate(AOVS):
+            prompt_embeds, _, pooled_prompt_embeds, _ = self._inverse.encode_prompt(
+                prompt=INVERSE_PROMPTS[aov_name],
+                prompt_2=None,
+                prompt_3=None,
+                do_classifier_free_guidance=False,
+            )
+            map_aware_mask = self.imaa(
+                patch_tokens=patch_tokens,
+                output_size=output_size,
+                map_ids=torch.tensor([map_index], device=image_tensor.device),
+            )
+            attn_mask = build_attn_mask(map_aware_mask, 154, img_len, 0.7)
+            result = self._inverse(
+                image=image_tensor,
+                prompt_embeds=prompt_embeds,
+                pooled_prompt_embeds=pooled_prompt_embeds,
+                guidance_scale=0.0,
+                image_guidance_scale=0.0,
+                num_inference_steps=num_inference_steps,
+                output_type=output_type,
+                aov=[aov_name],
+                map_aware_mask=attn_mask.to(image_tensor.device),
+            )
+            maps[aov_name] = result.images[0] if output_type == "pil" else result.images[0]
+        return maps
+
+    @torch.no_grad()
+    def render(
+        self,
+        maps: Dict[str, PipelineImageInput],
+        weather: str = "rainy",
+        num_inference_steps: int = 50,
+        guidance_scale: float = 6.0,
+        image_guidance_scale: float = 1.5,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        output_type: str = "pil",
+        render_size: Optional[int] = None,
+    ) -> Union[PIL.Image.Image, np.ndarray]:
+        r"""
+        Run forward weather rendering from intrinsic maps.
+
+        Args:
+            maps: Dict with keys `albedo`, `normal`, `roughness`, `metallic`, `irradiance`.
+            weather: Preset name (`rainy`, `sunny`, ...) or a custom prompt string.
+        """
+        if render_size is None:
+            sample = next(iter(maps.values()))
+            if isinstance(sample, PIL.Image.Image):
+                render_size = max(sample.size)
+            elif isinstance(sample, np.ndarray):
+                render_size = sample.shape[0]
+            else:
+                render_size = sample.shape[-1]
+
+        device = self._execution_device
+        dtype = self.dtype
+        aov_tensors = {
+            name: self._map_array_to_tensor(maps[name], render_size, device, dtype) for name in AOVS if name in maps
+        }
+        for name in AOVS:
+            if name not in aov_tensors:
+                aov_tensors[name] = torch.zeros((1, 3, render_size, render_size), device=device, dtype=dtype)
+
+        result = self._forward(
+            albedo=aov_tensors["albedo"],
+            normal=aov_tensors["normal"],
+            roughness=aov_tensors["roughness"],
+            metallic=aov_tensors["metallic"],
+            irradiance=aov_tensors.get("irradiance"),
+            prompt=[self._resolve_weather_prompt(weather)],
+            guidance_scale=guidance_scale,
+            image_guidance_scale=image_guidance_scale,
+            num_inference_steps=num_inference_steps,
+            required_aovs=AOVS,
+            generator=generator,
+            output_type=output_type,
+        )
+        return result.images[0]
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: PipelineImageInput,
+        weather: str = "rainy",
+        dino_model: Optional[PreTrainedModel] = None,
+        dino_processor: Optional[PreTrainedTokenizer] = None,
+        num_inverse_steps: int = 50,
+        num_forward_steps: int = 50,
+        guidance_scale: float = 6.0,
+        image_guidance_scale: float = 1.5,
+        image_size: Optional[int] = None,
+        render_size: Optional[int] = None,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        output_type: str = "pil",
+        return_maps: bool = False,
+        return_dict: bool = True,
+    ):
+        r"""
+        Decompose an RGB image into intrinsic maps, then render a weather-conditioned RGB image.
+
+        Examples:
+
+        Args:
+            image: Source RGB photograph or render.
+            weather: Weather preset (`rainy`, `sunny`, `snowy`, `foggy`, ...) or custom prompt.
+            dino_model (`transformers.PreTrainedModel`, *optional*):
+                DINOv2 model for IMAA gating. Required unless maps are passed via `maps=`.
+            dino_processor: DINO image processor paired with `dino_model`.
+            num_inverse_steps: Denoising steps for each intrinsic map.
+            num_forward_steps: Denoising steps for forward rendering.
+            guidance_scale: Text CFG for forward rendering.
+            image_guidance_scale: Image CFG for forward rendering.
+            image_size: Square size for inverse stage.
+            render_size: Square size for forward stage. Defaults to `image_size`.
+            return_maps: If `True`, also return decomposed intrinsic maps.
+            return_dict: Return [`IntrinsicWeatherPipelineOutput`] when `True`.
+        """
+        if dino_model is None or dino_processor is None:
+            raise ValueError("`dino_model` and `dino_processor` are required for end-to-end weather editing.")
+
+        maps = self.decompose(
+            image=image,
+            dino_model=dino_model,
+            dino_processor=dino_processor,
+            num_inference_steps=num_inverse_steps,
+            image_size=image_size,
+            output_type="np",
+        )
+        rendered = self.render(
+            maps=maps,
+            weather=weather,
+            num_inference_steps=num_forward_steps,
+            guidance_scale=guidance_scale,
+            image_guidance_scale=image_guidance_scale,
+            generator=generator,
+            output_type=output_type,
+            render_size=render_size or image_size,
+        )
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (rendered, maps) if return_maps else (rendered,)
+
+        return IntrinsicWeatherPipelineOutput(
+            images=[rendered] if output_type == "pil" else rendered,
+            maps=maps if return_maps else None,
+        )

pipeline_intrinsic_weather_forward.py

@@ -0,0 +1,1191 @@
+"""Hub custom pipeline: IntrinsicWeatherForwardPipeline.
+Load with native Hugging Face diffusers and trust_remote_code=True.
+"""
+
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import inspect
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import PIL.Image
+import torch
+import torch.nn.functional as F
+import torchvision
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import SD3LoraLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.models.transformers import SD3Transformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers
+from diffusers.utils import (
+    deprecate,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput
+
+from pipeline_utils import (
+    load_transformer_from_subfolder,
+    load_transformer_lora,
+    resolve_repo_dir,
+    resolve_transformer_lora_dir,
+    set_flow_timesteps,
+)
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from pipeline_intrinsic_weather_forward import IntrinsicWeatherForwardPipeline
+
+        >>> repo_dir = Path("./IntrisicWeather-diffusers").resolve()
+        >>> pipe = IntrinsicWeatherForwardPipeline.from_pretrained(
+        ...     repo_dir,
+        ...     transformer_subfolder="forward",
+        ...     local_files_only=True,
+        ...     torch_dtype=torch.float16,
+        ...     load_lora=True,
+        ... )
+        >>> pipe.to("cuda")
+        >>> image = pipe(
+        ...     prompt="A rainy day.",
+        ...     albedo=albedo_tensor,
+        ...     normal=normal_tensor,
+        ...     roughness=roughness_tensor,
+        ...     metallic=metallic_tensor,
+        ...     irradiance=irradiance_tensor,
+        ...     num_inference_steps=50,
+        ... ).images[0]
+        ```
+"""
+
+
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+    
+class VaeAoVProcessor(VaeImageProcessor):
+    def postprocess(
+        self,
+        image: torch.FloatTensor,
+        output_type: str = "pil",
+        do_denormalize: Optional[List[bool]] = None,
+        do_gamma_correction: bool = True,
+    ):
+        if not isinstance(image, torch.Tensor):
+            raise ValueError(
+                f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
+            )
+        if output_type not in ["latent", "pt", "np", "pil"]:
+            deprecation_message = (
+                f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
+                "`pil`, `np`, `pt`, `latent`"
+            )
+            deprecate(
+                "Unsupported output_type",
+                "1.0.0",
+                deprecation_message,
+                standard_warn=False,
+            )
+            output_type = "np"
+
+        if output_type == "latent":
+            return image
+
+        if do_denormalize is None:
+            do_denormalize = [self.config.do_normalize] * image.shape[0]
+
+        image = torch.stack(
+            [
+                self.denormalize(image[i]) if do_denormalize[i] else image[i]
+                for i in range(image.shape[0])
+            ]
+        )
+
+        # Gamma correction
+        if do_gamma_correction:
+            # image = torch.pow(image, 1.0 / 2.2)
+            image = image ** (1 / 2.2)
+
+        if output_type == "pt":
+            return image
+
+        image = self.pt_to_numpy(image)
+
+        if output_type == "np":
+            return image
+
+        if output_type == "pil":
+            return self.numpy_to_pil(image)
+
+class IntrinsicWeatherForwardPipeline(DiffusionPipeline, SD3LoraLoaderMixin):
+    r"""
+    Args:
+        transformer ([`SD3Transformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant,
+            with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size`
+            as its dimension.
+        text_encoder_2 ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the
+            [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
+            variant.
+        text_encoder_3 ([`T5EncoderModel`]):
+            Frozen text-encoder. Stable Diffusion 3 uses
+            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
+            [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`CLIPTokenizer`):
+            Second Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_3 (`T5TokenizerFast`):
+            Tokenizer of class
+            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"]
+
+    @classmethod
+    def load_transformer(
+        cls,
+        transformer_subfolder: str,
+        pretrained_model_name_or_path: str | Path,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        device: str | torch.device | None = None,
+    ):
+        return load_transformer_from_subfolder(
+            pretrained_model_name_or_path,
+            transformer_subfolder,
+            dtype=dtype,
+            device=device,
+        )
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        transformer_subfolder: str = "forward",
+        transformer=None,
+        load_lora: bool = True,
+        **kwargs,
+    ):
+        repo_dir = resolve_repo_dir(pretrained_model_name_or_path)
+        dtype = kwargs.get("torch_dtype", torch.bfloat16)
+        device = kwargs.get("device")
+        if transformer is None:
+            transformer = cls.load_transformer(
+                transformer_subfolder,
+                repo_dir,
+                dtype=dtype,
+                device=device,
+            )
+        pipe = DiffusionPipeline.from_pretrained(
+            repo_dir.as_posix(),
+            *args,
+            transformer=transformer,
+            custom_pipeline=str(repo_dir / "pipeline_intrinsic_weather_forward.py"),
+            trust_remote_code=True,
+            **kwargs,
+        )
+        lora_dir = resolve_transformer_lora_dir(repo_dir, transformer_subfolder)
+        if load_lora and lora_dir is not None:
+            pipe.load_lora_weights(lora_dir.as_posix())
+        if device is not None:
+            pipe = pipe.to(device)
+        return pipe
+
+    def __init__(
+        self,
+        transformer: SD3Transformer2DModel,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer_2: CLIPTokenizer,
+        text_encoder_3: T5EncoderModel,
+        tokenizer_3: T5TokenizerFast,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            text_encoder_3=text_encoder_3,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            tokenizer_3=tokenizer_3,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
+        )
+        self.image_processor = VaeAoVProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = (
+            self.transformer.config.sample_size
+            if hasattr(self, "transformer") and self.transformer is not None
+            else 128
+        )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if self.text_encoder_3 is None:
+            return torch.zeros(
+                (batch_size, self.tokenizer_max_length, self.transformer.config.joint_attention_dim),
+                device=device,
+                dtype=dtype,
+            )
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_3(text_input_ids.to(device))[0]
+
+        dtype = self.text_encoder_3.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        clip_skip: Optional[int] = None,
+        clip_model_index: int = 0,
+    ):
+        device = device or self._execution_device
+
+        clip_tokenizers = [self.tokenizer, self.tokenizer_2]
+        clip_text_encoders = [self.text_encoder, self.text_encoder_2]
+
+        tokenizer = clip_tokenizers[clip_model_index]
+        text_encoder = clip_text_encoders[clip_model_index]
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+        pooled_prompt_embeds = prompt_embeds[0]
+
+        if clip_skip is None:
+            prompt_embeds = prompt_embeds.hidden_states[-2]
+        else:
+            prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds, pooled_prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        prompt_3: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        clip_skip: Optional[int] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used in both text-encoders
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            prompt_3 = prompt_3 or prompt
+            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+            prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=0,
+            )
+            prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                prompt=prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=1,
+            )
+            clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
+
+            t5_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=prompt_3,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+            )
+
+            clip_prompt_embeds = torch.nn.functional.pad(
+                clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1])
+            )
+
+            prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            negative_prompt_3 = negative_prompt_3 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+            negative_prompt_3 = (
+                batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+            )
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds(
+                negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=0,
+            )
+            negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                negative_prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=1,
+            )
+            negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1)
+
+            t5_negative_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=negative_prompt_3, num_images_per_prompt=num_images_per_prompt, device=device
+            )
+
+            negative_clip_prompt_embeds = torch.nn.functional.pad(
+                negative_clip_prompt_embeds,
+                (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]),
+            )
+
+            negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2)
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
+            )
+
+        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        prompt_3,
+        # height,
+        # width,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        negative_prompt_3=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        # if height % 8 != 0 or width % 8 != 0:
+        #     raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_3 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_3`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
+            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_3 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        shape = (
+            batch_size,
+            num_channels_latents,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        return latents
+    
+    def prepare_image_latents(
+        self, image, batch_size, num_images_per_prompt, dtype, device, do_classifier_free_guidance, generator=None
+    ):
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        image = image.to(device=device, dtype=dtype)
+
+        batch_size = batch_size * num_images_per_prompt
+
+        if image.shape[1] == self.vae.config.latent_channels:
+            image_latents = image
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax")
+            # ? normalize image latents
+            # image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+        if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+            # expand image_latents for batch_size
+            deprecation_message = (
+                f"You have passed {batch_size} text prompts (`prompt`), but only {image_latents.shape[0]} initial"
+                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"
+                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"
+                " your script to pass as many initial images as text prompts to suppress this warning."
+            )
+            deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)
+            additional_image_per_prompt = batch_size // image_latents.shape[0]
+            image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+        elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            image_latents = torch.cat([image_latents], dim=0)
+
+        if do_classifier_free_guidance:
+            uncond_image_latents = torch.zeros_like(image_latents)
+            image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0)
+
+        return image_latents
+
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+    @property
+    def image_guidance_scale(self):
+        return self._image_guidance_scale
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        return self.guidance_scale > 1.0 and self.image_guidance_scale >= 1.0
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_3: Optional[Union[str, List[str]]] = None,
+        albedo: PipelineImageInput = None,
+        normal : PipelineImageInput = None,
+        metallic : PipelineImageInput = None,
+        roughness : PipelineImageInput = None,
+        irradiance : PipelineImageInput = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        timesteps: List[int] = None,
+        guidance_scale: float = 7.0,
+        image_guidance_scale: float = 1.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        mask_img: Optional[PipelineImageInput] = None,
+        required_aovs: List[str] = ["albedo"],
+        **kwargs
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used instead
+            negative_prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used instead
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        # height = height or self.default_sample_size * self.vae_scale_factor
+        # width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._image_guidance_scale = image_guidance_scale
+        self._clip_skip = clip_skip
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            device=device,
+            clip_skip=self.clip_skip,
+            num_images_per_prompt=num_images_per_prompt,
+        )
+        # print("prompt:", prompt_embeds.shape)
+        if self.do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds, negative_prompt_embeds], dim=0)
+            
+            # Similiarly
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, negative_pooled_prompt_embeds, negative_pooled_prompt_embeds], dim=0)          
+            
+        # if self.do_classifier_free_guidance:
+        #     prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+        #     pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+        # 3. Preprocess image
+        # Get image dimensions from the first available AOV (albedo, normal, roughness, metallic, irradiance)
+        old_height, old_width = None, None
+        aov_dict = {"albedo": albedo, "normal": normal, "roughness": roughness, "metallic": metallic, "irradiance": irradiance}
+        for aov_name in required_aovs:
+            aov_value = aov_dict.get(aov_name)
+            if aov_value is not None:
+                if isinstance(aov_value, PIL.Image.Image):
+                    old_width, old_height = aov_value.size
+                    break
+                elif isinstance(aov_value, torch.Tensor):
+                    if len(aov_value.shape) == 4:
+                        old_height = aov_value.shape[2]
+                        old_width = aov_value.shape[3]
+                    else:
+                        old_height = aov_value.shape[1]
+                        old_width = aov_value.shape[2]
+                    break
+        
+        # Fallback: if no AOV available, use default dimensions
+        if old_height is None or old_width is None:
+            old_height, old_width = 1024, 1024  # default fallback
+
+  
+        # For all aovs, the preprocessing remap the values to [-1, 1]
+        preprocessed_aovs = {}
+        for aov_name in required_aovs:
+            if aov_name == "albedo":
+                if albedo is not None:
+                    preprocessed_aovs[aov_name] = self.image_processor.preprocess(
+                        albedo
+                    )
+                else:
+                    preprocessed_aovs[aov_name] = None
+
+            if aov_name == "normal":
+                if normal is not None:
+                    preprocessed_aovs[aov_name] = (
+                        self.image_processor.preprocess(normal)
+                    )
+                else:
+                    preprocessed_aovs[aov_name] = None
+
+            if aov_name == "roughness":
+                if roughness is not None:
+                    preprocessed_aovs[aov_name] = self.image_processor.preprocess(
+                        roughness
+                    )
+                else:
+                    preprocessed_aovs[aov_name] = None
+            if aov_name == "metallic":
+                if metallic is not None:
+                    preprocessed_aovs[aov_name] = self.image_processor.preprocess(
+                        metallic
+                    )
+                else:
+                    preprocessed_aovs[aov_name] = None
+            if aov_name == "irradiance":
+                if irradiance is not None:
+                    preprocessed_aovs[aov_name] = self.image_processor.preprocess(
+                        irradiance
+                    )
+                else:
+                    preprocessed_aovs[aov_name] = None
+        # print("aovs:", preprocessed_aovs.keys())
+        # 4. Prepare latent variables
+        num_channels_latents = self.vae.config.latent_channels
+        # height, width = image_latents.shape[-2:]
+        height = old_height / 8
+        width = old_width / 8
+        height = height * self.vae_scale_factor
+        width = width * self.vae_scale_factor
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        height_latent, width_latent = latents.shape[-2:]
+
+        # 5. Prepare timesteps
+        set_flow_timesteps(
+            self.scheduler,
+            self.transformer,
+            num_inference_steps,
+            height_latent,
+            width_latent,
+            device,
+        )
+        timesteps = self.scheduler.timesteps
+
+        # timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # 6. Prepare Image latent
+        image_latents = []
+        for aov_name, aov in preprocessed_aovs.items():
+            # print(aov_name)
+            if aov is None:
+                image_latent = torch.zeros(
+                    batch_size,
+                    num_channels_latents,
+                    height_latent,
+                    width_latent,
+                    dtype=prompt_embeds.dtype,
+                    device=device,
+                )
+                # print(image_latent.shape)
+                if self.do_classifier_free_guidance:
+                    image_latents.append(
+                        torch.cat([image_latent, image_latent, image_latent], dim=0)
+                    )
+                else:
+                    image_latents.append(image_latent)
+            else:
+                image_latent = (
+                    self.prepare_image_latents(
+                        aov,
+                        batch_size,
+                        num_images_per_prompt,
+                        prompt_embeds.dtype,
+                        device,
+                        self.do_classifier_free_guidance,
+                        generator,
+                    )
+                )
+                image_latents.append(image_latent)
+        image_latents = torch.cat(image_latents, dim=1)
+
+        # 7. Check that shapes of latents and image match the DIT in_channels
+        num_channels_image = image_latents.shape[1]
+
+        if num_channels_latents + num_channels_image != self.transformer.config.in_channels:
+            raise ValueError(
+                f"Incorrect configuration settings! The config of `pipeline.transformer`: {self.transformer.config} expects"
+                f" {self.transformer.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
+                f" `num_channels_image`: {num_channels_image} "
+                f" = {num_channels_latents+num_channels_image}. Please verify the config of"
+                " `pipeline.transformer` or your `image` input."
+            )
+            
+        # 8. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 3) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+                
+                scaled_latent_model_input = torch.cat([latent_model_input, image_latents], dim=1)
+              
+                # if "mask_index" in kwargs and kwargs['mask_index'] is not None:
+                #     mask_index = kwargs['mask_index']
+                # else:
+                #     mask_index = None
+                noise_pred = self.transformer(
+                    hidden_states=scaled_latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    pooled_projections=pooled_prompt_embeds,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                    # mask_index= mask_index,
+                )[0]
+
+                # perform guidance 
+                if self.do_classifier_free_guidance:
+                    noise_pred_text, noise_pred_image, noise_pred_uncond = noise_pred.chunk(3)
+                    noise_pred = (
+                    noise_pred_uncond
+                    + self.guidance_scale * (noise_pred_text - noise_pred_image)
+                    + self.image_guidance_scale * (noise_pred_image - noise_pred_uncond)
+                )
+                    # noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) # neg, prompt
+                    # noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                    negative_pooled_prompt_embeds = callback_outputs.pop(
+                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+                    )
+                    image_latents = callback_outputs.pop("image_latents", image_latents)
+                    if mask_img is not None:
+                        mask_image_latents = callback_outputs.pop("mask_image_latents", mask_image_latents)
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            # latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            latents = latents / self.vae.config.scaling_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+            do_denormalize = [True] * image.shape[0]
+            
+            image = torchvision.transforms.Resize((old_height, old_width),interpolation=PIL.Image.BICUBIC)(image)
+            image = self.image_processor.postprocess(
+            image,
+            output_type=output_type,
+            do_denormalize=do_denormalize,
+            do_gamma_correction=False
+        )
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusion3PipelineOutput(images=image)

pipeline_intrinsic_weather_inverse.py

@@ -0,0 +1,1119 @@
+"""Hub custom pipeline: IntrinsicWeatherInversePipeline.
+Load with native Hugging Face diffusers and trust_remote_code=True.
+"""
+
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import inspect
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import PIL.Image
+import torch
+import torchvision
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import SD3LoraLoaderMixin
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler, KarrasDiffusionSchedulers
+from diffusers.utils import (
+    deprecate,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput
+
+from pipeline_utils import load_transformer_from_subfolder, resolve_repo_dir, set_flow_timesteps
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> from pathlib import Path
+        >>> import torch
+        >>> from pipeline_intrinsic_weather_inverse import IntrinsicWeatherInversePipeline
+
+        >>> repo_dir = Path("./IntrisicWeather-diffusers").resolve()
+        >>> transformer = IntrinsicWeatherInversePipeline.load_transformer("inverse-512", repo_dir)
+        >>> pipe = IntrinsicWeatherInversePipeline.from_pretrained(
+        ...     repo_dir,
+        ...     transformer_subfolder="inverse-512",
+        ...     local_files_only=True,
+        ...     torch_dtype=torch.float16,
+        ... )
+        >>> pipe.to("cuda")
+        >>> output = pipe(
+        ...     image=input_rgb_tensor,
+        ...     prompt="Albedo (diffuse basecolor)",
+        ...     aov=["albedo"],
+        ...     num_inference_steps=50,
+        ... )
+        ```
+"""
+
+
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+    
+class VaeAoVProcessor(VaeImageProcessor):
+    def postprocess(
+        self,
+        image: torch.FloatTensor,
+        output_type: str = "pil",
+        do_denormalize: Optional[List[bool]] = None,
+        do_gamma_correction: bool = True,
+    ):
+        if not isinstance(image, torch.Tensor):
+            raise ValueError(
+                f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
+            )
+        if output_type not in ["latent", "pt", "np", "pil"]:
+            deprecation_message = (
+                f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
+                "`pil`, `np`, `pt`, `latent`"
+            )
+            deprecate(
+                "Unsupported output_type",
+                "1.0.0",
+                deprecation_message,
+                standard_warn=False,
+            )
+            output_type = "np"
+
+        if output_type == "latent":
+            return image
+
+        if do_denormalize is None:
+            do_denormalize = [self.config.do_normalize] * image.shape[0]
+
+        image = torch.stack(
+            [
+                self.denormalize(image[i]) if do_denormalize[i] else image[i]
+                for i in range(image.shape[0])
+            ]
+        )
+
+        # Gamma correction
+        if do_gamma_correction:
+            # image = torch.pow(image, 1.0 / 2.2)
+            image = image ** (1 / 2.2)
+
+        if output_type == "pt":
+            return image
+
+        image = self.pt_to_numpy(image)
+
+        if output_type == "np":
+            return image
+
+        if output_type == "pil":
+            return self.numpy_to_pil(image)
+
+class IntrinsicWeatherInversePipeline(DiffusionPipeline, SD3LoraLoaderMixin):
+    r"""
+    Args:
+        transformer ([`SD3Transformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant,
+            with an additional added projection layer that is initialized with a diagonal matrix with the `hidden_size`
+            as its dimension.
+        text_encoder_2 ([`CLIPTextModelWithProjection`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+            specifically the
+            [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
+            variant.
+        text_encoder_3 ([`T5EncoderModel`]):
+            Frozen text-encoder. Stable Diffusion 3 uses
+            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
+            [t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`CLIPTokenizer`):
+            Second Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_3 (`T5TokenizerFast`):
+            Tokenizer of class
+            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"]
+
+    @classmethod
+    def load_transformer(
+        cls,
+        transformer_subfolder: str,
+        pretrained_model_name_or_path: str | Path,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        device: str | torch.device | None = None,
+    ):
+        return load_transformer_from_subfolder(
+            pretrained_model_name_or_path,
+            transformer_subfolder,
+            dtype=dtype,
+            device=device,
+        )
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *args,
+        transformer_subfolder: str = "inverse-512",
+        transformer=None,
+        **kwargs,
+    ):
+        repo_dir = resolve_repo_dir(pretrained_model_name_or_path)
+        dtype = kwargs.get("torch_dtype", torch.bfloat16)
+        device = kwargs.get("device")
+        if transformer is None:
+            transformer = cls.load_transformer(
+                transformer_subfolder,
+                repo_dir,
+                dtype=dtype,
+                device=device,
+            )
+        pipe = DiffusionPipeline.from_pretrained(
+            repo_dir.as_posix(),
+            *args,
+            transformer=transformer,
+            custom_pipeline=str(repo_dir / "pipeline_intrinsic_weather_inverse.py"),
+            trust_remote_code=True,
+            **kwargs,
+        )
+        if device is not None:
+            pipe = pipe.to(device)
+        return pipe
+
+    def __init__(
+        self,
+        transformer,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer_2: CLIPTokenizer,
+        text_encoder_3: T5EncoderModel,
+        tokenizer_3: T5TokenizerFast,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            text_encoder_3=text_encoder_3,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            tokenizer_3=tokenizer_3,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
+        )
+        self.image_processor = VaeAoVProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = (
+            self.transformer.config.sample_size
+            if hasattr(self, "transformer") and self.transformer is not None
+            else 128
+        )
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if self.text_encoder_3 is None:
+            return torch.zeros(
+                (batch_size, self.tokenizer_max_length, self.transformer.config.joint_attention_dim),
+                device=device,
+                dtype=dtype,
+            )
+
+        text_inputs = self.tokenizer_3(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_3(text_input_ids.to(device))[0]
+
+        dtype = self.text_encoder_3.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        clip_skip: Optional[int] = None,
+        clip_model_index: int = 0,
+    ):
+        device = device or self._execution_device
+
+        clip_tokenizers = [self.tokenizer, self.tokenizer_2]
+        clip_text_encoders = [self.text_encoder, self.text_encoder_2]
+
+        tokenizer = clip_tokenizers[clip_model_index]
+        text_encoder = clip_text_encoders[clip_model_index]
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+        pooled_prompt_embeds = prompt_embeds[0]
+
+        if clip_skip is None:
+            prompt_embeds = prompt_embeds.hidden_states[-2]
+        else:
+            prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds, pooled_prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]] = None,
+        prompt_3: Union[str, List[str]] = None,
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        clip_skip: Optional[int] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used in both text-encoders
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            prompt_3 = prompt_3 or prompt
+            prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3
+
+            prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=0,
+            )
+            prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                prompt=prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=clip_skip,
+                clip_model_index=1,
+            )
+            clip_prompt_embeds = torch.cat([prompt_embed, prompt_2_embed], dim=-1)
+
+            t5_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=prompt_3,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+            )
+
+            clip_prompt_embeds = torch.nn.functional.pad(
+                clip_prompt_embeds, (0, t5_prompt_embed.shape[-1] - clip_prompt_embeds.shape[-1])
+            )
+
+            prompt_embeds = torch.cat([clip_prompt_embeds, t5_prompt_embed], dim=-2)
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embed, pooled_prompt_2_embed], dim=-1)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+            negative_prompt_3 = negative_prompt_3 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
+            negative_prompt_3 = (
+                batch_size * [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3
+            )
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds(
+                negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=0,
+            )
+            negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds(
+                negative_prompt_2,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                clip_skip=None,
+                clip_model_index=1,
+            )
+            negative_clip_prompt_embeds = torch.cat([negative_prompt_embed, negative_prompt_2_embed], dim=-1)
+
+            t5_negative_prompt_embed = self._get_t5_prompt_embeds(
+                prompt=negative_prompt_3, num_images_per_prompt=num_images_per_prompt, device=device
+            )
+
+            negative_clip_prompt_embeds = torch.nn.functional.pad(
+                negative_clip_prompt_embeds,
+                (0, t5_negative_prompt_embed.shape[-1] - negative_clip_prompt_embeds.shape[-1]),
+            )
+
+            negative_prompt_embeds = torch.cat([negative_clip_prompt_embeds, t5_negative_prompt_embed], dim=-2)
+            negative_pooled_prompt_embeds = torch.cat(
+                [negative_pooled_prompt_embed, negative_pooled_prompt_2_embed], dim=-1
+            )
+
+        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        prompt_3,
+        # height,
+        # width,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        negative_prompt_3=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        # if height % 8 != 0 or width % 8 != 0:
+        #     raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_3 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_3`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+        elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)):
+            raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_3 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        shape = (
+            batch_size,
+            num_channels_latents,
+            int(height) // self.vae_scale_factor,
+            int(width) // self.vae_scale_factor,
+        )
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        return latents
+    
+    def prepare_image_latents(
+        self, image, batch_size, num_images_per_prompt, dtype, device, do_classifier_free_guidance, generator=None
+    ):
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        image = image.to(device=device, dtype=dtype)
+
+        batch_size = batch_size * num_images_per_prompt
+
+        if image.shape[1] == self.vae.config.latent_channels:
+            image_latents = image
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax")
+            # ? normalize image latents
+            # image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+        if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+            # expand image_latents for batch_size
+            deprecation_message = (
+                f"You have passed {batch_size} text prompts (`prompt`), but only {image_latents.shape[0]} initial"
+                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"
+                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"
+                " your script to pass as many initial images as text prompts to suppress this warning."
+            )
+            deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)
+            additional_image_per_prompt = batch_size // image_latents.shape[0]
+            image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+        elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            image_latents = torch.cat([image_latents], dim=0)
+
+        if do_classifier_free_guidance:
+            uncond_image_latents = torch.zeros_like(image_latents)
+            image_latents = torch.cat([image_latents, image_latents, uncond_image_latents], dim=0)
+
+        return image_latents
+
+
+
+
+
+
+
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+    @property
+    def image_guidance_scale(self):
+        return self._image_guidance_scale
+
+    @property
+    def clip_skip(self):
+        return self._clip_skip
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        return self.guidance_scale > 1.0 and self.image_guidance_scale >= 1.0
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        prompt_3: Optional[Union[str, List[str]]] = None,
+        image: PipelineImageInput = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        timesteps: List[int] = None,
+        guidance_scale: float = 7.0,
+        image_guidance_scale: float = 1.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_3: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        mask_img: Optional[PipelineImageInput] = None,
+        aov: List[str] = ["albedo"],
+        map_aware_mask: Optional[torch.FloatTensor] = None,
+        **kwargs
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 5.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used instead
+            negative_prompt_3 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and
+                `text_encoder_3`. If not defined, `negative_prompt` is used instead
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        # height = height or self.default_sample_size * self.vae_scale_factor
+        # width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            prompt_3,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._image_guidance_scale = image_guidance_scale
+        self._clip_skip = clip_skip
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_3=prompt_3,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            negative_prompt_3=negative_prompt_3,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            device=device,
+            clip_skip=self.clip_skip,
+            num_images_per_prompt=num_images_per_prompt,
+        )
+        # print("prompt:", prompt_embeds.shape)
+        if self.do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds, negative_prompt_embeds], dim=0)
+            
+            # Similiarly
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, negative_pooled_prompt_embeds, negative_pooled_prompt_embeds], dim=0)          
+            
+        # if self.do_classifier_free_guidance:
+        #     prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+        #     pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)
+
+        # 3. Preprocess image
+        if (isinstance(image, PIL.Image.Image) and len(image.size)==4) or (isinstance(image, torch.Tensor) and len(image.shape) == 4):
+            old_height = image.shape[2]
+            old_width = image.shape[3]
+        else:
+            old_height = image.shape[1]
+            old_width = image.shape[2]
+
+        image = self.image_processor.preprocess(image)
+        # image = torchvision.transforms.Resize((512,512),interpolation=PIL.Image.BICUBIC)(image)
+
+        # 4. Prepare Image latent
+        image_latents = self.prepare_image_latents(
+            image,
+            batch_size,
+            num_images_per_prompt,
+            prompt_embeds.dtype,
+            device,
+            self.do_classifier_free_guidance,
+        )
+
+        # 5. Prepare timesteps
+        set_flow_timesteps(
+            self.scheduler,
+            self.transformer,
+            num_inference_steps,
+            image_latents.shape[-2],
+            image_latents.shape[-1],
+            device,
+        )
+        timesteps = self.scheduler.timesteps
+
+        # timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        height, width = image_latents.shape[-2:]
+        height = height * self.vae_scale_factor
+        width = width * self.vae_scale_factor
+        # 6. Prepare latent variables
+        num_channels_latents = self.vae.config.latent_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 7. Check that shapes of latents and image match the DIT in_channels
+        num_channels_image = image_latents.shape[1]
+        if mask_img is not None:
+            mask_img = self.image_processor.preprocess(mask_img)
+            mask_image_latents = self.prepare_image_latents(
+                mask_img,
+                batch_size,
+                num_images_per_prompt,
+                prompt_embeds.dtype,
+                device,
+                self.do_classifier_free_guidance,
+            )
+            num_channels_image += mask_image_latents.shape[1]
+
+        if num_channels_latents + num_channels_image != self.transformer.config.in_channels:
+            raise ValueError(
+                f"Incorrect configuration settings! The config of `pipeline.transformer`: {self.transformer.config} expects"
+                f" {self.transformer.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
+                f" `num_channels_image`: {num_channels_image} "
+                f" = {num_channels_latents+num_channels_image}. Please verify the config of"
+                " `pipeline.transformer` or your `image` input."
+            )
+            
+        # 8. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 3) if self.do_classifier_free_guidance else latents
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+                
+                scaled_latent_model_input = torch.cat([latent_model_input, image_latents], dim=1)
+                if mask_img is not None:
+                    scaled_latent_model_input = torch.cat([scaled_latent_model_input, mask_image_latents], dim=1)
+                # if "mask_index" in kwargs and kwargs['mask_index'] is not None:
+                #     mask_index = kwargs['mask_index']
+                # else:
+                #     mask_index = None
+                noise_pred = self.transformer(
+                    hidden_states=scaled_latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    pooled_projections=pooled_prompt_embeds,
+                    map_aware_mask=map_aware_mask,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                    # mask_index= mask_index,
+                )[0]
+
+                # perform guidance 
+                if self.do_classifier_free_guidance:
+                    noise_pred_text, noise_pred_image, noise_pred_uncond = noise_pred.chunk(3)
+                    noise_pred = (
+                    noise_pred_uncond
+                    + self.guidance_scale * (noise_pred_text - noise_pred_image)
+                    + self.image_guidance_scale * (noise_pred_image - noise_pred_uncond)
+                )
+                    # noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) # neg, prompt
+                    # noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                    negative_pooled_prompt_embeds = callback_outputs.pop(
+                        "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+                    )
+                    image_latents = callback_outputs.pop("image_latents", image_latents)
+                    if mask_img is not None:
+                        mask_image_latents = callback_outputs.pop("mask_image_latents", mask_image_latents)
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            # latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            latents = latents / self.vae.config.scaling_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+
+            do_denormalize = [True] * image.shape[0]
+            aov_name = aov[0]
+            if aov_name == "albedo" or aov_name == "irradiance":
+                do_gamma_correction = True
+            else:
+                do_gamma_correction = False
+        
+            if aov_name == "roughness" or aov_name == "metallic":
+                image = image[:, 0:1].repeat(1, 3, 1, 1)
+            # print(image.shape)
+            # print(old_height, old_width)
+            image = torchvision.transforms.Resize((old_height, old_width),interpolation=PIL.Image.BICUBIC)(image)
+            image = self.image_processor.postprocess(
+            image,
+            output_type=output_type,
+            do_denormalize=do_denormalize,
+            do_gamma_correction=do_gamma_correction,
+        )
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return StableDiffusion3PipelineOutput(images=image)

pipeline_utils.py

@@ -0,0 +1,104 @@
+"""Helpers for loading transformer variants from ``transformer/<subfolder>/``."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import torch
+from diffusers.models.transformers import SD3Transformer2DModel
+
+
+def calculate_shift(
+    image_seq_len: int,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+) -> float:
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    return image_seq_len * m + b
+
+
+def set_flow_timesteps(
+    scheduler,
+    transformer,
+    num_inference_steps: int,
+    latent_height: int,
+    latent_width: int,
+    device: torch.device,
+) -> None:
+    if scheduler.config.get("use_dynamic_shifting", False):
+        image_seq_len = (latent_height // transformer.config.patch_size) * (
+            latent_width // transformer.config.patch_size
+        )
+        mu = calculate_shift(
+            image_seq_len,
+            scheduler.config.get("base_image_seq_len", 256),
+            scheduler.config.get("max_image_seq_len", 4096),
+            scheduler.config.get("base_shift", 0.5),
+            scheduler.config.get("max_shift", 1.15),
+        )
+        scheduler.set_timesteps(num_inference_steps, device=device, mu=mu)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device)
+
+
+def resolve_repo_dir(pretrained_model_name_or_path: str | Path) -> Path:
+    return Path(pretrained_model_name_or_path).resolve()
+
+
+def load_transformer_from_subfolder(
+    repo_dir: str | Path,
+    transformer_subfolder: str,
+    *,
+    dtype: torch.dtype = torch.bfloat16,
+    device: str | torch.device | None = None,
+):
+    """Load a transformer checkpoint from ``<repo_dir>/transformer/<transformer_subfolder>/``."""
+    repo_dir = resolve_repo_dir(repo_dir)
+    transformer_path = repo_dir / "transformer" / transformer_subfolder
+    if not transformer_path.is_dir():
+        raise FileNotFoundError(f"Transformer folder not found: {transformer_path}")
+
+    custom_module = transformer_path / "transformer_intrinsic_weather.py"
+    if custom_module.exists():
+        spec = importlib.util.spec_from_file_location("transformer_intrinsic_weather", custom_module)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Cannot import custom transformer module: {custom_module}")
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        cls = module.IntrinsicWeatherSD3Transformer2DModel
+        transformer = cls.from_pretrained(
+            transformer_path.as_posix(),
+            torch_dtype=dtype,
+            local_files_only=True,
+        )
+    else:
+        transformer = SD3Transformer2DModel.from_pretrained(
+            transformer_path.as_posix(),
+            torch_dtype=dtype,
+            local_files_only=True,
+        )
+
+    if device is not None:
+        transformer = transformer.to(device)
+    return transformer
+
+
+def resolve_transformer_lora_dir(repo_dir: str | Path, transformer_subfolder: str) -> Path | None:
+    """Return ``transformer/<subfolder>/lora`` when present."""
+    lora_dir = resolve_repo_dir(repo_dir) / "transformer" / transformer_subfolder / "lora"
+    if lora_dir.is_dir() and any(lora_dir.glob("*.safetensors")):
+        return lora_dir
+    return None
+
+
+def load_transformer_lora(pipe, repo_dir: str | Path, transformer_subfolder: str) -> bool:
+    """Load LoRA weights bundled with a transformer variant. Returns True if loaded."""
+    lora_dir = resolve_transformer_lora_dir(repo_dir, transformer_subfolder)
+    if lora_dir is None:
+        return False
+    pipe.load_lora_weights(lora_dir.as_posix())
+    return True

scheduler/scheduler_config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.35.1",
+  "num_train_timesteps": 1000,
+  "shift": 3.0,
+  "use_dynamic_shifting": true,
+  "base_shift": 0.5,
+  "max_shift": 1.15
+}

test_all_pipelines.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Smoke-test all IntrinsicWeather pipelines on CUDA with bfloat16."""
+
+from __future__ import annotations
+
+import gc
+import sys
+from pathlib import Path
+
+import torch
+
+REPO = Path(__file__).resolve().parent
+sys.path.insert(0, str(REPO))
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+IMAGE_SIZE = 512
+STEPS = 2
+
+
+def _clear():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def test_inverse() -> None:
+    from imaa.imaa import IMAA
+    from pipeline_intrinsic_weather_inverse import IntrinsicWeatherInversePipeline
+    from safetensors.torch import load_file
+
+    print("[inverse] loading pipeline ...")
+    pipe = IntrinsicWeatherInversePipeline.from_pretrained(
+        REPO,
+        transformer_subfolder="inverse-512",
+        device=DEVICE,
+        local_files_only=True,
+        torch_dtype=DTYPE,
+    )
+    assert next(pipe.transformer.parameters()).dtype == DTYPE
+    assert next(pipe.transformer.parameters()).device.type == "cuda"
+    print(f"[inverse] transformer in_channels={pipe.transformer.config.in_channels}")
+
+    imaa = IMAA(dino_model=None, processor=None, num_maps=5, map_embedding_dim=256, common_dim=128).to(DEVICE)
+    imaa.load_state_dict(load_file((REPO / "imaa" / "model.safetensors").as_posix()))
+    imaa.eval()
+
+    image = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=DEVICE, dtype=DTYPE)
+    prompt_embeds, _, pooled_prompt_embeds, _ = pipe.encode_prompt(
+        prompt="Albedo (diffuse basecolor)",
+        prompt_2=None,
+        prompt_3=None,
+        do_classifier_free_guidance=False,
+    )
+    print("[inverse] running 2-step inference ...")
+    out = pipe(
+        image=image,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        guidance_scale=0.0,
+        image_guidance_scale=0.0,
+        num_inference_steps=STEPS,
+        output_type="pt",
+        aov=["albedo"],
+        map_aware_mask=None,
+    )
+    print(f"[inverse] output shape={tuple(out.images[0].shape)} dtype={out.images[0].dtype}")
+    del pipe, imaa, out
+    _clear()
+
+
+def test_forward() -> None:
+    from pipeline_intrinsic_weather_forward import IntrinsicWeatherForwardPipeline
+
+    print("[forward] loading pipeline ...")
+    pipe = IntrinsicWeatherForwardPipeline.from_pretrained(
+        REPO,
+        transformer_subfolder="forward",
+        device=DEVICE,
+        local_files_only=True,
+        torch_dtype=DTYPE,
+        load_lora=True,
+    )
+    assert next(pipe.transformer.parameters()).dtype == DTYPE
+    assert next(pipe.transformer.parameters()).device.type == "cuda"
+    print(f"[forward] transformer in_channels={pipe.transformer.config.in_channels}")
+
+    aov = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=DEVICE, dtype=DTYPE)
+    print("[forward] running 2-step inference ...")
+    out = pipe(
+        albedo=aov,
+        normal=aov,
+        roughness=aov,
+        metallic=aov,
+        irradiance=aov,
+        prompt=["A rainy day."],
+        guidance_scale=6.0,
+        image_guidance_scale=1.5,
+        num_inference_steps=STEPS,
+        required_aovs=["albedo", "normal", "roughness", "metallic", "irradiance"],
+        generator=torch.Generator(device=DEVICE).manual_seed(0),
+    )
+    print(f"[forward] output size={out.images[0].size}")
+    del pipe, out
+    _clear()
+
+
+def test_unified() -> None:
+    from pipeline_intrinsic_weather import IntrinsicWeatherPipeline
+
+    print("[unified] loading pipeline ...")
+    pipe = IntrinsicWeatherPipeline.from_pretrained(
+        REPO,
+        inverse_transformer_subfolder="inverse-512",
+        forward_transformer_subfolder="forward",
+        device=DEVICE,
+        local_files_only=True,
+        torch_dtype=DTYPE,
+        load_lora=True,
+        load_imaa=True,
+    )
+    assert next(pipe.inverse_transformer.parameters()).dtype == DTYPE
+    assert next(pipe.forward_transformer.parameters()).dtype == DTYPE
+    print("[unified] loaded inverse + forward transformers and IMAA")
+    del pipe
+    _clear()
+
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for this test.")
+
+    print(f"device={torch.cuda.get_device_name(0)} dtype={DTYPE}")
+    test_inverse()
+    test_forward()
+    test_unified()
+    print("All pipeline tests passed.")
+
+
+if __name__ == "__main__":
+    main()

text_encoder/config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "CLIPTextModelWithProjection"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 768,
+  "torch_dtype": "float16",
+  "transformers_version": "4.41.2",
+  "vocab_size": 49408
+}

text_encoder/model.safetensors

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

text_encoder_2/config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "CLIPTextModelWithProjection"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_size": 1280,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 20,
+  "num_hidden_layers": 32,
+  "pad_token_id": 1,
+  "projection_dim": 1280,
+  "torch_dtype": "float16",
+  "transformers_version": "4.41.2",
+  "vocab_size": 49408
+}

text_encoder_2/model.safetensors

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

text_encoder_3/config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "T5EncoderModel"
+  ],
+  "classifier_dropout": 0.0,
+  "d_ff": 10240,
+  "d_kv": 64,
+  "d_model": 4096,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "gelu_new",
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "num_decoder_layers": 24,
+  "num_heads": 64,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float16",
+  "transformers_version": "4.41.2",
+  "use_cache": true,
+  "vocab_size": 32128
+}

text_encoder_3/model-00001-of-00002.safetensors

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

text_encoder_3/model-00002-of-00002.safetensors

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

text_encoder_3/model.safetensors.index.json

@@ -0,0 +1,226 @@
+{
+  "metadata": {
+    "total_size": 9524621312
+  },
+  "weight_map": {
+    "encoder.block.0.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.12.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.13.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.14.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.15.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.16.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.17.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.18.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.19.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.20.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.21.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.22.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.23.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.4.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.5.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.6.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.7.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.8.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.9.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.final_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "shared.weight": "model-00001-of-00002.safetensors"
+  }
+}

tokenizer/special_tokens_map.json

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

tokenizer/tokenizer_config.json

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

tokenizer_2/special_tokens_map.json

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

tokenizer_2/tokenizer_config.json

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

tokenizer_3/special_tokens_map.json

@@ -0,0 +1,125 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_3/spiece.model

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

tokenizer_3/tokenizer_config.json

@@ -0,0 +1,940 @@
+{
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<extra_id_99>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<extra_id_98>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<extra_id_97>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<extra_id_96>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<extra_id_95>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<extra_id_94>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<extra_id_93>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<extra_id_92>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<extra_id_91>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32009": {
+      "content": "<extra_id_90>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32010": {
+      "content": "<extra_id_89>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32011": {
+      "content": "<extra_id_88>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32012": {
+      "content": "<extra_id_87>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32013": {
+      "content": "<extra_id_86>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32014": {
+      "content": "<extra_id_85>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32015": {
+      "content": "<extra_id_84>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32016": {
+      "content": "<extra_id_83>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32017": {
+      "content": "<extra_id_82>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32018": {
+      "content": "<extra_id_81>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32019": {
+      "content": "<extra_id_80>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32020": {
+      "content": "<extra_id_79>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32021": {
+      "content": "<extra_id_78>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32022": {
+      "content": "<extra_id_77>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32023": {
+      "content": "<extra_id_76>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32024": {
+      "content": "<extra_id_75>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32025": {
+      "content": "<extra_id_74>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32026": {
+      "content": "<extra_id_73>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32027": {
+      "content": "<extra_id_72>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32028": {
+      "content": "<extra_id_71>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32029": {
+      "content": "<extra_id_70>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32030": {
+      "content": "<extra_id_69>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32031": {
+      "content": "<extra_id_68>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32032": {
+      "content": "<extra_id_67>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32033": {
+      "content": "<extra_id_66>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32034": {
+      "content": "<extra_id_65>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32035": {
+      "content": "<extra_id_64>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32036": {
+      "content": "<extra_id_63>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32037": {
+      "content": "<extra_id_62>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32038": {
+      "content": "<extra_id_61>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32039": {
+      "content": "<extra_id_60>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32040": {
+      "content": "<extra_id_59>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32041": {
+      "content": "<extra_id_58>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32042": {
+      "content": "<extra_id_57>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32043": {
+      "content": "<extra_id_56>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32044": {
+      "content": "<extra_id_55>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32045": {
+      "content": "<extra_id_54>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32046": {
+      "content": "<extra_id_53>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32047": {
+      "content": "<extra_id_52>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32048": {
+      "content": "<extra_id_51>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32049": {
+      "content": "<extra_id_50>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32050": {
+      "content": "<extra_id_49>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32051": {
+      "content": "<extra_id_48>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32052": {
+      "content": "<extra_id_47>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32053": {
+      "content": "<extra_id_46>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32054": {
+      "content": "<extra_id_45>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32055": {
+      "content": "<extra_id_44>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32056": {
+      "content": "<extra_id_43>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32057": {
+      "content": "<extra_id_42>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32058": {
+      "content": "<extra_id_41>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32059": {
+      "content": "<extra_id_40>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32060": {
+      "content": "<extra_id_39>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32061": {
+      "content": "<extra_id_38>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32062": {
+      "content": "<extra_id_37>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32063": {
+      "content": "<extra_id_36>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32064": {
+      "content": "<extra_id_35>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32065": {
+      "content": "<extra_id_34>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32066": {
+      "content": "<extra_id_33>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32067": {
+      "content": "<extra_id_32>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32068": {
+      "content": "<extra_id_31>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32069": {
+      "content": "<extra_id_30>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32070": {
+      "content": "<extra_id_29>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32071": {
+      "content": "<extra_id_28>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32072": {
+      "content": "<extra_id_27>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32073": {
+      "content": "<extra_id_26>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32074": {
+      "content": "<extra_id_25>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32075": {
+      "content": "<extra_id_24>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32076": {
+      "content": "<extra_id_23>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32077": {
+      "content": "<extra_id_22>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32078": {
+      "content": "<extra_id_21>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32079": {
+      "content": "<extra_id_20>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32080": {
+      "content": "<extra_id_19>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32081": {
+      "content": "<extra_id_18>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32082": {
+      "content": "<extra_id_17>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32083": {
+      "content": "<extra_id_16>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32084": {
+      "content": "<extra_id_15>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32085": {
+      "content": "<extra_id_14>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32086": {
+      "content": "<extra_id_13>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32087": {
+      "content": "<extra_id_12>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32088": {
+      "content": "<extra_id_11>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32089": {
+      "content": "<extra_id_10>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32090": {
+      "content": "<extra_id_9>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32091": {
+      "content": "<extra_id_8>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32092": {
+      "content": "<extra_id_7>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32093": {
+      "content": "<extra_id_6>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32094": {
+      "content": "<extra_id_5>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32095": {
+      "content": "<extra_id_4>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32096": {
+      "content": "<extra_id_3>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32097": {
+      "content": "<extra_id_2>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32098": {
+      "content": "<extra_id_1>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32099": {
+      "content": "<extra_id_0>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "legacy": true,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

transformer/forward/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "SD3Transformer2DModel",
+  "_diffusers_version": "0.38.0",
+  "attention_head_dim": 64,
+  "caption_projection_dim": 1536,
+  "dual_attention_layers": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12
+  ],
+  "in_channels": 96,
+  "joint_attention_dim": 4096,
+  "num_attention_heads": 24,
+  "num_layers": 24,
+  "out_channels": 16,
+  "patch_size": 2,
+  "pooled_projection_dim": 2048,
+  "pos_embed_max_size": 384,
+  "qk_norm": "rms_norm",
+  "sample_size": 128
+}

transformer/forward/diffusion_pytorch_model.safetensors

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

transformer/forward/lora/pytorch_lora_weights.safetensors

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

transformer/inverse-512/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "IntrinsicWeatherSD3Transformer2DModel",
+  "_diffusers_version": "0.38.0",
+  "attention_head_dim": 64,
+  "caption_projection_dim": 1536,
+  "dual_attention_layers": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12
+  ],
+  "in_channels": 32,
+  "joint_attention_dim": 4096,
+  "num_attention_heads": 24,
+  "num_layers": 24,
+  "out_channels": 16,
+  "patch_size": 2,
+  "pooled_projection_dim": 2048,
+  "pos_embed_max_size": 384,
+  "qk_norm": "rms_norm",
+  "sample_size": 128
+}

transformer/inverse-512/diffusion_pytorch_model.safetensors

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

transformer/inverse-512/transformer_intrinsic_weather.py

@@ -0,0 +1,1527 @@
+# Copyright 2024 Stability AI, The HuggingFace Team and The InstantX Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import nn as torch_nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin, SD3Transformer2DLoadersMixin
+from diffusers.models.attention import FeedForward, JointTransformerBlock, _chunked_feed_forward
+from diffusers.models.attention_processor import (
+    Attention,
+    AttentionProcessor,
+    AttnProcessor,
+    AttnProcessor2_0,
+    FusedJointAttnProcessor2_0,
+    JointAttnProcessor2_0,
+    SpatialNorm,
+)
+from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed, SinusoidalPositionalEmbedding
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import (
+    AdaLayerNorm,
+    AdaLayerNormContinuous,
+    AdaLayerNormZero,
+    RMSNorm,
+    SD35AdaLayerNormZeroX,
+)
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+
+logger = logging.get_logger(__name__)
+
+
+class MapAwareAttention(nn.Module):
+    r"""
+    A cross attention layer.
+
+    Parameters:
+        query_dim (`int`):
+            The number of channels in the query.
+        cross_attention_dim (`int`, *optional*):
+            The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.
+        heads (`int`,  *optional*, defaults to 8):
+            The number of heads to use for multi-head attention.
+        kv_heads (`int`,  *optional*, defaults to `None`):
+            The number of key and value heads to use for multi-head attention. Defaults to `heads`. If
+            `kv_heads=heads`, the model will use Multi Head Attention (MHA), if `kv_heads=1` the model will use Multi
+            Query Attention (MQA) otherwise GQA is used.
+        dim_head (`int`,  *optional*, defaults to 64):
+            The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0):
+            The dropout probability to use.
+        bias (`bool`, *optional*, defaults to False):
+            Set to `True` for the query, key, and value linear layers to contain a bias parameter.
+        upcast_attention (`bool`, *optional*, defaults to False):
+            Set to `True` to upcast the attention computation to `float32`.
+        upcast_softmax (`bool`, *optional*, defaults to False):
+            Set to `True` to upcast the softmax computation to `float32`.
+        cross_attention_norm (`str`, *optional*, defaults to `None`):
+            The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
+        cross_attention_norm_num_groups (`int`, *optional*, defaults to 32):
+            The number of groups to use for the group norm in the cross attention.
+        added_kv_proj_dim (`int`, *optional*, defaults to `None`):
+            The number of channels to use for the added key and value projections. If `None`, no projection is used.
+        norm_num_groups (`int`, *optional*, defaults to `None`):
+            The number of groups to use for the group norm in the attention.
+        spatial_norm_dim (`int`, *optional*, defaults to `None`):
+            The number of channels to use for the spatial normalization.
+        out_bias (`bool`, *optional*, defaults to `True`):
+            Set to `True` to use a bias in the output linear layer.
+        scale_qk (`bool`, *optional*, defaults to `True`):
+            Set to `True` to scale the query and key by `1 / sqrt(dim_head)`.
+        only_cross_attention (`bool`, *optional*, defaults to `False`):
+            Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if
+            `added_kv_proj_dim` is not `None`.
+        eps (`float`, *optional*, defaults to 1e-5):
+            An additional value added to the denominator in group normalization that is used for numerical stability.
+        rescale_output_factor (`float`, *optional*, defaults to 1.0):
+            A factor to rescale the output by dividing it with this value.
+        residual_connection (`bool`, *optional*, defaults to `False`):
+            Set to `True` to add the residual connection to the output.
+        _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`):
+            Set to `True` if the attention block is loaded from a deprecated state dict.
+        processor (`AttnProcessor`, *optional*, defaults to `None`):
+            The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and
+            `AttnProcessor` otherwise.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: Optional[int] = None,
+        heads: int = 8,
+        kv_heads: Optional[int] = None,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        upcast_attention: bool = False,
+        upcast_softmax: bool = False,
+        cross_attention_norm: Optional[str] = None,
+        cross_attention_norm_num_groups: int = 32,
+        qk_norm: Optional[str] = None,
+        added_kv_proj_dim: Optional[int] = None,
+        added_proj_bias: Optional[bool] = True,
+        norm_num_groups: Optional[int] = None,
+        spatial_norm_dim: Optional[int] = None,
+        out_bias: bool = True,
+        scale_qk: bool = True,
+        only_cross_attention: bool = False,
+        eps: float = 1e-5,
+        rescale_output_factor: float = 1.0,
+        residual_connection: bool = False,
+        _from_deprecated_attn_block: bool = False,
+        processor: Optional["AttnProcessor"] = None,
+        out_dim: int = None,
+        out_context_dim: int = None,
+        context_pre_only=None,
+        pre_only=False,
+        elementwise_affine: bool = True,
+        is_causal: bool = False,
+    ):
+        super().__init__()
+
+        # To prevent circular import.
+        from diffusers.models.normalization import FP32LayerNorm, LpNorm, RMSNorm
+
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+        self.query_dim = query_dim
+        self.use_bias = bias
+        self.is_cross_attention = cross_attention_dim is not None
+        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+        self.upcast_attention = upcast_attention
+        self.upcast_softmax = upcast_softmax
+        self.rescale_output_factor = rescale_output_factor
+        self.residual_connection = residual_connection
+        self.dropout = dropout
+        self.fused_projections = False
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
+        self.context_pre_only = context_pre_only
+        self.pre_only = pre_only
+        self.is_causal = is_causal
+
+        # we make use of this private variable to know whether this class is loaded
+        # with an deprecated state dict so that we can convert it on the fly
+        self._from_deprecated_attn_block = _from_deprecated_attn_block
+
+        self.scale_qk = scale_qk
+        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+        self.heads = out_dim // dim_head if out_dim is not None else heads
+        # for slice_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set slice_size with `set_attention_slice`
+        self.sliceable_head_dim = heads
+
+        self.added_kv_proj_dim = added_kv_proj_dim
+        self.only_cross_attention = only_cross_attention
+
+        if self.added_kv_proj_dim is None and self.only_cross_attention:
+            raise ValueError(
+                "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
+            )
+
+        if norm_num_groups is not None:
+            self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)
+        else:
+            self.group_norm = None
+
+        if spatial_norm_dim is not None:
+            self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim)
+        else:
+            self.spatial_norm = None
+
+        if qk_norm is None:
+            self.norm_q = None
+            self.norm_k = None
+        elif qk_norm == "layer_norm":
+            self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+            self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+        elif qk_norm == "fp32_layer_norm":
+            self.norm_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+            self.norm_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+        elif qk_norm == "layer_norm_across_heads":
+            # Lumina applies qk norm across all heads
+            self.norm_q = nn.LayerNorm(dim_head * heads, eps=eps)
+            self.norm_k = nn.LayerNorm(dim_head * kv_heads, eps=eps)
+        elif qk_norm == "rms_norm":
+            self.norm_q = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+            self.norm_k = RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+        elif qk_norm == "rms_norm_across_heads":
+            # LTX applies qk norm across all heads
+            self.norm_q = RMSNorm(dim_head * heads, eps=eps)
+            self.norm_k = RMSNorm(dim_head * kv_heads, eps=eps)
+        elif qk_norm == "l2":
+            self.norm_q = LpNorm(p=2, dim=-1, eps=eps)
+            self.norm_k = LpNorm(p=2, dim=-1, eps=eps)
+        else:
+            raise ValueError(
+                f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'."
+            )
+
+        if cross_attention_norm is None:
+            self.norm_cross = None
+        elif cross_attention_norm == "layer_norm":
+            self.norm_cross = nn.LayerNorm(self.cross_attention_dim)
+        elif cross_attention_norm == "group_norm":
+            if self.added_kv_proj_dim is not None:
+                # The given `encoder_hidden_states` are initially of shape
+                # (batch_size, seq_len, added_kv_proj_dim) before being projected
+                # to (batch_size, seq_len, cross_attention_dim). The norm is applied
+                # before the projection, so we need to use `added_kv_proj_dim` as
+                # the number of channels for the group norm.
+                norm_cross_num_channels = added_kv_proj_dim
+            else:
+                norm_cross_num_channels = self.cross_attention_dim
+
+            self.norm_cross = nn.GroupNorm(
+                num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True
+            )
+        else:
+            raise ValueError(
+                f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
+            )
+
+        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+        if not self.only_cross_attention:
+            # only relevant for the `AddedKVProcessor` classes
+            self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+            self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+        else:
+            self.to_k = None
+            self.to_v = None
+
+        self.added_proj_bias = added_proj_bias
+        if self.added_kv_proj_dim is not None:
+            self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+            self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+            if self.context_pre_only is not None:
+                self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+        else:
+            self.add_q_proj = None
+            self.add_k_proj = None
+            self.add_v_proj = None
+
+        if not self.pre_only:
+            self.to_out = nn.ModuleList([])
+            self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+            self.to_out.append(nn.Dropout(dropout))
+        else:
+            self.to_out = None
+
+        if self.context_pre_only is not None and not self.context_pre_only:
+            self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+        else:
+            self.to_add_out = None
+
+        if qk_norm is not None and added_kv_proj_dim is not None:
+            if qk_norm == "layer_norm":
+                self.norm_added_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+                self.norm_added_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+            elif qk_norm == "fp32_layer_norm":
+                self.norm_added_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+                self.norm_added_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+            elif qk_norm == "rms_norm":
+                self.norm_added_q = RMSNorm(dim_head, eps=eps)
+                self.norm_added_k = RMSNorm(dim_head, eps=eps)
+            elif qk_norm == "rms_norm_across_heads":
+                # Wan applies qk norm across all heads
+                # Wan also doesn't apply a q norm
+                self.norm_added_q = None
+                self.norm_added_k = RMSNorm(dim_head * kv_heads, eps=eps)
+            else:
+                raise ValueError(
+                    f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`"
+                )
+        else:
+            self.norm_added_q = None
+            self.norm_added_k = None
+
+        # set attention processor
+        # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+        # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+        # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+        if processor is None:
+            processor = (
+                AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+            )
+        self.set_processor(processor)
+
+    # def set_use_xla_flash_attention(
+    #     self,
+    #     use_xla_flash_attention: bool,
+    #     partition_spec: Optional[Tuple[Optional[str], ...]] = None,
+    #     is_flux=False,
+    # ) -> None:
+    #     r"""
+    #     Set whether to use xla flash attention from `torch_xla` or not.
+
+    #     Args:
+    #         use_xla_flash_attention (`bool`):
+    #             Whether to use pallas flash attention kernel from `torch_xla` or not.
+    #         partition_spec (`Tuple[]`, *optional*):
+    #             Specify the partition specification if using SPMD. Otherwise None.
+    #     """
+    #     if use_xla_flash_attention:
+    #         if not is_torch_xla_available:
+    #             raise "torch_xla is not available"
+    #         elif is_torch_xla_version("<", "2.3"):
+    #             raise "flash attention pallas kernel is supported from torch_xla version 2.3"
+    #         elif is_spmd() and is_torch_xla_version("<", "2.4"):
+    #             raise "flash attention pallas kernel using SPMD is supported from torch_xla version 2.4"
+    #         else:
+    #             if is_flux:
+    #                 processor = XLAFluxFlashAttnProcessor2_0(partition_spec)
+    #             else:
+    #                 processor = XLAFlashAttnProcessor2_0(partition_spec)
+    #     else:
+    #         processor = (
+    #             AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+    #         )
+    #     self.set_processor(processor)
+
+    # def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None:
+    #     r"""
+    #     Set whether to use npu flash attention from `torch_npu` or not.
+
+    #     """
+    #     if use_npu_flash_attention:
+    #         processor = AttnProcessorNPU()
+    #     else:
+    #         # set attention processor
+    #         # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+    #         # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+    #         # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+    #         processor = (
+    #             AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+    #         )
+    #     self.set_processor(processor)
+
+    # def set_use_memory_efficient_attention_xformers(
+    #     self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
+    # ) -> None:
+    #     r"""
+    #     Set whether to use memory efficient attention from `xformers` or not.
+
+    #     Args:
+    #         use_memory_efficient_attention_xformers (`bool`):
+    #             Whether to use memory efficient attention from `xformers` or not.
+    #         attention_op (`Callable`, *optional*):
+    #             The attention operation to use. Defaults to `None` which uses the default attention operation from
+    #             `xformers`.
+    #     """
+    #     is_custom_diffusion = hasattr(self, "processor") and isinstance(
+    #         self.processor,
+    #         (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor, CustomDiffusionAttnProcessor2_0),
+    #     )
+    #     is_added_kv_processor = hasattr(self, "processor") and isinstance(
+    #         self.processor,
+    #         (
+    #             AttnAddedKVProcessor,
+    #             AttnAddedKVProcessor2_0,
+    #             SlicedAttnAddedKVProcessor,
+    #             XFormersAttnAddedKVProcessor,
+    #         ),
+    #     )
+    #     is_ip_adapter = hasattr(self, "processor") and isinstance(
+    #         self.processor,
+    #         (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0, IPAdapterXFormersAttnProcessor),
+    #     )
+    #     is_joint_processor = hasattr(self, "processor") and isinstance(
+    #         self.processor,
+    #         (
+    #             JointAttnProcessor2_0,
+    #             XFormersJointAttnProcessor,
+    #         ),
+    #     )
+
+    #     if use_memory_efficient_attention_xformers:
+    #         if is_added_kv_processor and is_custom_diffusion:
+    #             raise NotImplementedError(
+    #                 f"Memory efficient attention is currently not supported for custom diffusion for attention processor type {self.processor}"
+    #             )
+    #         if not is_xformers_available():
+    #             raise ModuleNotFoundError(
+    #                 (
+    #                     "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
+    #                     " xformers"
+    #                 ),
+    #                 name="xformers",
+    #             )
+    #         elif not torch.cuda.is_available():
+    #             raise ValueError(
+    #                 "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"
+    #                 " only available for GPU "
+    #             )
+    #         else:
+    #             try:
+    #                 # Make sure we can run the memory efficient attention
+    #                 dtype = None
+    #                 if attention_op is not None:
+    #                     op_fw, op_bw = attention_op
+    #                     dtype, *_ = op_fw.SUPPORTED_DTYPES
+    #                 q = torch.randn((1, 2, 40), device="cuda", dtype=dtype)
+    #                 _ = xformers.ops.memory_efficient_attention(q, q, q)
+    #             except Exception as e:
+    #                 raise e
+
+    #         if is_custom_diffusion:
+    #             processor = CustomDiffusionXFormersAttnProcessor(
+    #                 train_kv=self.processor.train_kv,
+    #                 train_q_out=self.processor.train_q_out,
+    #                 hidden_size=self.processor.hidden_size,
+    #                 cross_attention_dim=self.processor.cross_attention_dim,
+    #                 attention_op=attention_op,
+    #             )
+    #             processor.load_state_dict(self.processor.state_dict())
+    #             if hasattr(self.processor, "to_k_custom_diffusion"):
+    #                 processor.to(self.processor.to_k_custom_diffusion.weight.device)
+    #         elif is_added_kv_processor:
+    #             # TODO(Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP
+    #             # which uses this type of cross attention ONLY because the attention mask of format
+    #             # [0, ..., -10.000, ..., 0, ...,] is not supported
+    #             # throw warning
+    #             logger.info(
+    #                 "Memory efficient attention with `xformers` might currently not work correctly if an attention mask is required for the attention operation."
+    #             )
+    #             processor = XFormersAttnAddedKVProcessor(attention_op=attention_op)
+    #         elif is_ip_adapter:
+    #             processor = IPAdapterXFormersAttnProcessor(
+    #                 hidden_size=self.processor.hidden_size,
+    #                 cross_attention_dim=self.processor.cross_attention_dim,
+    #                 num_tokens=self.processor.num_tokens,
+    #                 scale=self.processor.scale,
+    #                 attention_op=attention_op,
+    #             )
+    #             processor.load_state_dict(self.processor.state_dict())
+    #             if hasattr(self.processor, "to_k_ip"):
+    #                 processor.to(
+    #                     device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype
+    #                 )
+    #         elif is_joint_processor:
+    #             processor = XFormersJointAttnProcessor(attention_op=attention_op)
+    #         else:
+    #             processor = XFormersAttnProcessor(attention_op=attention_op)
+    #     else:
+    #         if is_custom_diffusion:
+    #             attn_processor_class = (
+    #                 CustomDiffusionAttnProcessor2_0
+    #                 if hasattr(F, "scaled_dot_product_attention")
+    #                 else CustomDiffusionAttnProcessor
+    #             )
+    #             processor = attn_processor_class(
+    #                 train_kv=self.processor.train_kv,
+    #                 train_q_out=self.processor.train_q_out,
+    #                 hidden_size=self.processor.hidden_size,
+    #                 cross_attention_dim=self.processor.cross_attention_dim,
+    #             )
+    #             processor.load_state_dict(self.processor.state_dict())
+    #             if hasattr(self.processor, "to_k_custom_diffusion"):
+    #                 processor.to(self.processor.to_k_custom_diffusion.weight.device)
+    #         elif is_ip_adapter:
+    #             processor = IPAdapterAttnProcessor2_0(
+    #                 hidden_size=self.processor.hidden_size,
+    #                 cross_attention_dim=self.processor.cross_attention_dim,
+    #                 num_tokens=self.processor.num_tokens,
+    #                 scale=self.processor.scale,
+    #             )
+    #             processor.load_state_dict(self.processor.state_dict())
+    #             if hasattr(self.processor, "to_k_ip"):
+    #                 processor.to(
+    #                     device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype
+    #                 )
+    #         else:
+    #             # set attention processor
+    #             # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+    #             # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+    #             # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+    #             processor = (
+    #                 AttnProcessor2_0()
+    #                 if hasattr(F, "scaled_dot_product_attention") and self.scale_qk
+    #                 else AttnProcessor()
+    #             )
+
+    #     self.set_processor(processor)
+
+    # def set_attention_slice(self, slice_size: int) -> None:
+    #     r"""
+    #     Set the slice size for attention computation.
+
+    #     Args:
+    #         slice_size (`int`):
+    #             The slice size for attention computation.
+    #     """
+    #     if slice_size is not None and slice_size > self.sliceable_head_dim:
+    #         raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.")
+
+    #     if slice_size is not None and self.added_kv_proj_dim is not None:
+    #         processor = SlicedAttnAddedKVProcessor(slice_size)
+    #     elif slice_size is not None:
+    #         processor = SlicedAttnProcessor(slice_size)
+    #     elif self.added_kv_proj_dim is not None:
+    #         processor = AttnAddedKVProcessor()
+    #     else:
+    #         # set attention processor
+    #         # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+    #         # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+    #         # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+    #         processor = (
+    #             AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+    #         )
+
+    #     self.set_processor(processor)
+
+    def set_processor(self, processor: "AttnProcessor") -> None:
+        r"""
+        Set the attention processor to use.
+
+        Args:
+            processor (`AttnProcessor`):
+                The attention processor to use.
+        """
+        # if current processor is in `self._modules` and if passed `processor` is not, we need to
+        # pop `processor` from `self._modules`
+        if (
+            hasattr(self, "processor")
+            and isinstance(self.processor, torch.nn.Module)
+            and not isinstance(processor, torch.nn.Module)
+        ):
+            logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}")
+            self._modules.pop("processor")
+
+        self.processor = processor
+
+    def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProcessor":
+        r"""
+        Get the attention processor in use.
+
+        Args:
+            return_deprecated_lora (`bool`, *optional*, defaults to `False`):
+                Set to `True` to return the deprecated LoRA attention processor.
+
+        Returns:
+            "AttentionProcessor": The attention processor in use.
+        """
+        if not return_deprecated_lora:
+            return self.processor
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        map_aware_mask: Optional[torch.FloatTensor] = None,
+        **cross_attention_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        The forward method of the `Attention` class.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                The hidden states of the query.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                The hidden states of the encoder.
+            map_aware_mask (`torch.Tensor`, *optional*):
+                The attention mask to use. If `None`, no mask is applied.
+            **cross_attention_kwargs:
+                Additional keyword arguments to pass along to the cross attention.
+
+        Returns:
+            `torch.Tensor`: The output of the attention layer.
+        """
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**cross_attention_kwargs` is empty
+
+        attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+        quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
+        unused_kwargs = [
+            k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters
+        ]
+        if len(unused_kwargs) > 0:
+            logger.warning(
+                f"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+            )
+        cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}
+
+        return self.processor(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=map_aware_mask,
+            **cross_attention_kwargs,
+        )
+
+    def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:
+        r"""
+        Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads`
+        is the number of heads initialized while constructing the `Attention` class.
+
+        Args:
+            tensor (`torch.Tensor`): The tensor to reshape.
+
+        Returns:
+            `torch.Tensor`: The reshaped tensor.
+        """
+        head_size = self.heads
+        batch_size, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+        tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+        return tensor
+
+    def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:
+        r"""
+        Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is
+        the number of heads initialized while constructing the `Attention` class.
+
+        Args:
+            tensor (`torch.Tensor`): The tensor to reshape.
+            out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is
+                reshaped to `[batch_size * heads, seq_len, dim // heads]`.
+
+        Returns:
+            `torch.Tensor`: The reshaped tensor.
+        """
+        head_size = self.heads
+        if tensor.ndim == 3:
+            batch_size, seq_len, dim = tensor.shape
+            extra_dim = 1
+        else:
+            batch_size, extra_dim, seq_len, dim = tensor.shape
+        tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size)
+        tensor = tensor.permute(0, 2, 1, 3)
+
+        if out_dim == 3:
+            tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size)
+
+        return tensor
+
+    def get_attention_scores(
+        self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        r"""
+        Compute the attention scores.
+
+        Args:
+            query (`torch.Tensor`): The query tensor.
+            key (`torch.Tensor`): The key tensor.
+            attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
+
+        Returns:
+            `torch.Tensor`: The attention probabilities/scores.
+        """
+        dtype = query.dtype
+        if self.upcast_attention:
+            query = query.float()
+            key = key.float()
+
+        if attention_mask is None:
+            baddbmm_input = torch.empty(
+                query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device
+            )
+            beta = 0
+        else:
+            baddbmm_input = attention_mask
+            beta = 1
+
+        attention_scores = torch.baddbmm(
+            baddbmm_input,
+            query,
+            key.transpose(-1, -2),
+            beta=beta,
+            alpha=self.scale,
+        )
+        del baddbmm_input
+
+        if self.upcast_softmax:
+            attention_scores = attention_scores.float()
+
+        attention_probs = attention_scores.softmax(dim=-1)
+        del attention_scores
+
+        attention_probs = attention_probs.to(dtype)
+
+        return attention_probs
+
+    def prepare_attention_mask(
+        self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3
+    ) -> torch.Tensor:
+        r"""
+        Prepare the attention mask for the attention computation.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                The attention mask to prepare.
+            target_length (`int`):
+                The target length of the attention mask. This is the length of the attention mask after padding.
+            batch_size (`int`):
+                The batch size, which is used to repeat the attention mask.
+            out_dim (`int`, *optional*, defaults to `3`):
+                The output dimension of the attention mask. Can be either `3` or `4`.
+
+        Returns:
+            `torch.Tensor`: The prepared attention mask.
+        """
+        head_size = self.heads
+        if attention_mask is None:
+            return attention_mask
+
+        current_length: int = attention_mask.shape[-1]
+        if current_length != target_length:
+            if attention_mask.device.type == "mps":
+                # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+                # Instead, we can manually construct the padding tensor.
+                padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)
+                padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)
+                attention_mask = torch.cat([attention_mask, padding], dim=2)
+            else:
+                # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+                #       we want to instead pad by (0, remaining_length), where remaining_length is:
+                #       remaining_length: int = target_length - current_length
+                # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+                attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
+
+        if out_dim == 3:
+            if attention_mask.shape[0] < batch_size * head_size:
+                attention_mask = attention_mask.repeat_interleave(
+                    head_size, dim=0, output_size=attention_mask.shape[0] * head_size
+                )
+        elif out_dim == 4:
+            attention_mask = attention_mask.unsqueeze(1)
+            attention_mask = attention_mask.repeat_interleave(
+                head_size, dim=1, output_size=attention_mask.shape[1] * head_size
+            )
+
+        return attention_mask
+
+    def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        r"""
+        Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the
+        `Attention` class.
+
+        Args:
+            encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder.
+
+        Returns:
+            `torch.Tensor`: The normalized encoder hidden states.
+        """
+        assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
+
+        if isinstance(self.norm_cross, nn.LayerNorm):
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+        elif isinstance(self.norm_cross, nn.GroupNorm):
+            # Group norm norms along the channels dimension and expects
+            # input to be in the shape of (N, C, *). In this case, we want
+            # to norm along the hidden dimension, so we need to move
+            # (batch_size, sequence_length, hidden_size) ->
+            # (batch_size, hidden_size, sequence_length)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+            encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+        else:
+            assert False
+
+        return encoder_hidden_states
+
+    @torch.no_grad()
+    def fuse_projections(self, fuse=True):
+        device = self.to_q.weight.data.device
+        dtype = self.to_q.weight.data.dtype
+
+        if not self.is_cross_attention:
+            # fetch weight matrices.
+            concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])
+            in_features = concatenated_weights.shape[1]
+            out_features = concatenated_weights.shape[0]
+
+            # create a new single projection layer and copy over the weights.
+            self.to_qkv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)
+            self.to_qkv.weight.copy_(concatenated_weights)
+            if self.use_bias:
+                concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])
+                self.to_qkv.bias.copy_(concatenated_bias)
+
+        else:
+            concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data])
+            in_features = concatenated_weights.shape[1]
+            out_features = concatenated_weights.shape[0]
+
+            self.to_kv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)
+            self.to_kv.weight.copy_(concatenated_weights)
+            if self.use_bias:
+                concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data])
+                self.to_kv.bias.copy_(concatenated_bias)
+
+        # handle added projections for SD3 and others.
+        if (
+            getattr(self, "add_q_proj", None) is not None
+            and getattr(self, "add_k_proj", None) is not None
+            and getattr(self, "add_v_proj", None) is not None
+        ):
+            concatenated_weights = torch.cat(
+                [self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data]
+            )
+            in_features = concatenated_weights.shape[1]
+            out_features = concatenated_weights.shape[0]
+
+            self.to_added_qkv = nn.Linear(
+                in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype
+            )
+            self.to_added_qkv.weight.copy_(concatenated_weights)
+            if self.added_proj_bias:
+                concatenated_bias = torch.cat(
+                    [self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data]
+                )
+                self.to_added_qkv.bias.copy_(concatenated_bias)
+
+        self.fused_projections = fuse
+
+class MapAwareAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("JointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        # print("attention_mask: ", attention_mask)
+        residual = hidden_states
+
+        batch_size = hidden_states.shape[0]
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # `context` projections.
+        if encoder_hidden_states is not None:
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # print(f"image q: {query.shape}, image k: {key.shape}, image v: {value.shape}")  # [B, 24, 1024, 64]
+            # print(f"text q: {encoder_hidden_states_query_proj.shape}, text k: {encoder_hidden_states_key_proj.shape}, text v: {encoder_hidden_states_value_proj.shape}")
+            # [B, 24, 154, 64]
+
+            query = torch.cat([query, encoder_hidden_states_query_proj], dim=2)
+            key = torch.cat([key, encoder_hidden_states_key_proj], dim=2)
+            value = torch.cat([value, encoder_hidden_states_value_proj], dim=2)
+
+            # print(f"Joint - query shape: {query.shape}, key shape: {key.shape}, value shape: {value.shape}")
+            # [B, 24, 1178, 64]
+            map_aware_mask = attention_mask
+        else:
+            map_aware_mask = None
+        # print(
+        #     "map_aware_mask:",
+        #     None if map_aware_mask is None else (map_aware_mask.shape, map_aware_mask.dtype)
+        # )
+
+        # print("query: ", query.shape, query.dtype)
+        if map_aware_mask is not None:
+            map_aware_mask = map_aware_mask.to(query.dtype)
+    
+
+        hidden_states = F.scaled_dot_product_attention(query, key, value, 
+                                                       attn_mask=map_aware_mask, 
+                                                       dropout_p=0.0, 
+                                                       is_causal=False)
+
+        
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            # Split the attention outputs.
+            hidden_states, encoder_hidden_states = (
+                hidden_states[:, : residual.shape[1]],
+                hidden_states[:, residual.shape[1] :],
+            )
+            if not attn.context_pre_only:
+                encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if encoder_hidden_states is not None:
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+class MapAwareTransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://huggingface.co/papers/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        context_pre_only: bool = False,
+        qk_norm: Optional[str] = None,
+        use_dual_attention: bool = False,
+    ):
+        super().__init__()
+
+        self.use_dual_attention = use_dual_attention
+        self.context_pre_only = context_pre_only
+        context_norm_type = "ada_norm_continous" if context_pre_only else "ada_norm_zero"
+
+        if use_dual_attention:
+            self.norm1 = SD35AdaLayerNormZeroX(dim)
+        else:
+            self.norm1 = AdaLayerNormZero(dim)
+
+        if context_norm_type == "ada_norm_continous":
+            self.norm1_context = AdaLayerNormContinuous(
+                dim, dim, elementwise_affine=False, eps=1e-6, bias=True, norm_type="layer_norm"
+            )
+        elif context_norm_type == "ada_norm_zero":
+            self.norm1_context = AdaLayerNormZero(dim)
+        else:
+            raise ValueError(
+                f"Unknown context_norm_type: {context_norm_type}, currently only support `ada_norm_continous`, `ada_norm_zero`"
+            )
+
+        # if hasattr(F, "scaled_dot_product_attention"):
+        #     processor = JointAttnProcessor2_0()
+        # else:
+        #     raise ValueError(
+        #         "The current PyTorch version does not support the `scaled_dot_product_attention` function."
+        #     )
+
+        self.attn = MapAwareAttention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=context_pre_only,
+            bias=True,
+            processor=MapAwareAttnProcessor2_0(),
+            qk_norm=qk_norm,
+            eps=1e-6,
+        )
+
+        if use_dual_attention:
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=None,
+                dim_head=attention_head_dim,
+                heads=num_attention_heads,
+                out_dim=dim,
+                bias=True,
+                processor=JointAttnProcessor2_0(),
+                qk_norm=qk_norm,
+                eps=1e-6,
+            )
+        else:
+            self.attn2 = None
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        if not context_pre_only:
+            self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+            self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        else:
+            self.norm2_context = None
+            self.ff_context = None
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    # Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        map_aware_mask: Optional[torch.FloatTensor] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        if self.use_dual_attention:
+            # print(f"hidden_states: {type(hidden_states)}")
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1(
+                hidden_states, emb=temb
+            )
+        else:
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+
+        if self.context_pre_only:
+            norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)
+        else:
+            norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+                encoder_hidden_states, emb=temb
+            )
+
+        # Attention.
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            map_aware_mask=map_aware_mask,
+            **joint_attention_kwargs,
+        )
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        if self.use_dual_attention:
+            attn_output2 = self.attn2(hidden_states=norm_hidden_states2, **joint_attention_kwargs)
+            attn_output2 = gate_msa2.unsqueeze(1) * attn_output2
+            hidden_states = hidden_states + attn_output2
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+        else:
+            ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = hidden_states + ff_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+        if self.context_pre_only:
+            encoder_hidden_states = None
+        else:
+            context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+            encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+            norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+            norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+            if self._chunk_size is not None:
+                # "feed_forward_chunk_size" can be used to save memory
+                context_ff_output = _chunked_feed_forward(
+                    self.ff_context, norm_encoder_hidden_states, self._chunk_dim, self._chunk_size
+                )
+            else:
+                context_ff_output = self.ff_context(norm_encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+
+        return encoder_hidden_states, hidden_states
+
+@maybe_allow_in_graph
+class SD3SingleTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+    ):
+        super().__init__()
+
+        self.norm1 = AdaLayerNormZero(dim)
+        self.attn = Attention(
+            query_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=JointAttnProcessor2_0(),
+            eps=1e-6,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor):
+        # 1. Attention
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+        attn_output = self.attn(hidden_states=norm_hidden_states, encoder_hidden_states=None)
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        # 2. Feed Forward
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+        hidden_states = hidden_states + ff_output
+
+        return hidden_states
+
+
+class IntrinsicWeatherSD3Transformer2DModel(
+    ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, SD3Transformer2DLoadersMixin
+):
+    """
+    The Transformer model introduced in [Stable Diffusion 3](https://huggingface.co/papers/2403.03206).
+
+    Parameters:
+        sample_size (`int`, defaults to `128`):
+            The width/height of the latents. This is fixed during training since it is used to learn a number of
+            position embeddings.
+        patch_size (`int`, defaults to `2`):
+            Patch size to turn the input data into small patches.
+        in_channels (`int`, defaults to `16`):
+            The number of latent channels in the input.
+        num_layers (`int`, defaults to `18`):
+            The number of layers of transformer blocks to use.
+        attention_head_dim (`int`, defaults to `64`):
+            The number of channels in each head.
+        num_attention_heads (`int`, defaults to `18`):
+            The number of heads to use for multi-head attention.
+        joint_attention_dim (`int`, defaults to `4096`):
+            The embedding dimension to use for joint text-image attention.
+        caption_projection_dim (`int`, defaults to `1152`):
+            The embedding dimension of caption embeddings.
+        pooled_projection_dim (`int`, defaults to `2048`):
+            The embedding dimension of pooled text projections.
+        out_channels (`int`, defaults to `16`):
+            The number of latent channels in the output.
+        pos_embed_max_size (`int`, defaults to `96`):
+            The maximum latent height/width of positional embeddings.
+        dual_attention_layers (`Tuple[int, ...]`, defaults to `()`):
+            The number of dual-stream transformer blocks to use.
+        qk_norm (`str`, *optional*, defaults to `None`):
+            The normalization to use for query and key in the attention layer. If `None`, no normalization is used.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["JointTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 128,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        num_layers: int = 18,
+        attention_head_dim: int = 64,
+        num_attention_heads: int = 18,
+        joint_attention_dim: int = 4096,
+        caption_projection_dim: int = 1152,
+        pooled_projection_dim: int = 2048,
+        out_channels: int = 16,
+        pos_embed_max_size: int = 96,
+        dual_attention_layers: Tuple[
+            int, ...
+        ] = (),  # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5
+        qk_norm: Optional[str] = None,
+    ):
+        super().__init__()
+        self.out_channels = out_channels if out_channels is not None else in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.pos_embed = PatchEmbed(
+            height=sample_size,
+            width=sample_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            embed_dim=self.inner_dim,
+            pos_embed_max_size=pos_embed_max_size,  # hard-code for now.
+        )
+        self.time_text_embed = CombinedTimestepTextProjEmbeddings(
+            embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+        )
+        self.context_embedder = nn.Linear(joint_attention_dim, caption_projection_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MapAwareTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    context_pre_only=i == num_layers - 1,
+                    qk_norm=qk_norm,
+                    use_dual_attention=True if i in dual_attention_layers else False,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+    # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
+    def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+        """
+        Sets the attention processor to use [feed forward
+        chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
+
+        Parameters:
+            chunk_size (`int`, *optional*):
+                The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
+                over each tensor of dim=`dim`.
+            dim (`int`, *optional*, defaults to `0`):
+                The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
+                or dim=1 (sequence length).
+        """
+        if dim not in [0, 1]:
+            raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
+
+        # By default chunk size is 1
+        chunk_size = chunk_size or 1
+
+        def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
+            if hasattr(module, "set_chunk_feed_forward"):
+                module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
+
+            for child in module.children():
+                fn_recursive_feed_forward(child, chunk_size, dim)
+
+        for module in self.children():
+            fn_recursive_feed_forward(module, chunk_size, dim)
+
+    # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking
+    def disable_forward_chunking(self):
+        def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
+            if hasattr(module, "set_chunk_feed_forward"):
+                module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
+
+            for child in module.children():
+                fn_recursive_feed_forward(child, chunk_size, dim)
+
+        for module in self.children():
+            fn_recursive_feed_forward(module, None, 0)
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedJointAttnProcessor2_0
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedJointAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        block_controlnet_hidden_states: List = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+        skip_layers: Optional[List[int]] = None,
+        map_aware_mask: Optional[torch.FloatTensor] = None,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        """
+        The [`SD3Transformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch size, channel, height, width)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch size, sequence_len, embed_dims)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`):
+                Embeddings projected from the embeddings of input conditions.
+            timestep (`torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+            skip_layers (`list` of `int`, *optional*):
+                A list of layer indices to skip during the forward pass.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        height, width = hidden_states.shape[-2:]
+
+        hidden_states = self.pos_embed(hidden_states)  # takes care of adding positional embeddings too.
+        temb = self.time_text_embed(timestep, pooled_projections)
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+            ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+            ip_hidden_states, ip_temb = self.image_proj(ip_adapter_image_embeds, timestep)
+
+            joint_attention_kwargs.update(ip_hidden_states=ip_hidden_states, temb=ip_temb)
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            # print("index: ", index_block)
+
+            # Skip specified layers
+            is_skip = True if skip_layers is not None and index_block in skip_layers else False
+
+            if index_block >= (self.config.num_layers // 2) and map_aware_mask is not None:
+                current_mask = map_aware_mask.to(hidden_states.device)
+            else:
+                current_mask = None
+
+            # print("transformer: map_aware_mask:", current_mask.shape if current_mask is not None else None)
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing and not is_skip:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    current_mask,
+                    joint_attention_kwargs,
+                )
+            elif not is_skip:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    map_aware_mask=current_mask,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if block_controlnet_hidden_states is not None and block.context_pre_only is False:
+                interval_control = len(self.transformer_blocks) / len(block_controlnet_hidden_states)
+                hidden_states = hidden_states + block_controlnet_hidden_states[int(index_block / interval_control)]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        # unpatchify
+        patch_size = self.config.patch_size
+        height = height // patch_size
+        width = width // patch_size
+
+        hidden_states = hidden_states.reshape(
+            shape=(hidden_states.shape[0], height, width, patch_size, patch_size, self.out_channels)
+        )
+        hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+        output = hidden_states.reshape(
+            shape=(hidden_states.shape[0], self.out_channels, height * patch_size, width * patch_size)
+        )
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

vae/config.json

@@ -0,0 +1,36 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.29.0.dev0",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 16,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 1.5305,
+  "shift_factor": 0.0609,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": false,
+  "use_quant_conv": false
+}

vae/diffusion_pytorch_model.safetensors

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