diff --git a/.pylintrc b/.pylintrc
index 2c9bbc0c2..5ecbbef48 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -27,6 +27,7 @@ ignore-paths=/usr/lib/.*$,
              modules/meissonic,
              modules/mod,
              modules/omnigen,
+             modules/omnigen2,
              modules/onnx_impl,
              modules/pag,
              modules/pixelsmith,
diff --git a/.ruff.toml b/.ruff.toml
index ab5e0601c..023678c33 100644
--- a/.ruff.toml
+++ b/.ruff.toml
@@ -20,6 +20,7 @@ exclude = [
     "modules/meissonic",
     "modules/mod",
     "modules/omnigen",
+    "modules/omnigen2",
     "modules/hidream",
     "modules/pag",
     "modules/pixelsmith",
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ed2ebb762..bc7e9f782 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,29 +1,120 @@
 # Change Log for SD.Next
 
+## Update for 2025-06-30
+
+### Highlights for 2025-06-30
+
+New release with ~100 commits...So what's new? Well, its been a busy few weeks with new models coming out quite frequently:  
+- New T2I/I2I models: **OmniGen-2, Cosmos-Predict2, FLUX.1-Kontext, Chroma**  
+- Additional VLM models: **JoyCaption Beta, MoonDream 2**  
+- Additional upscalers: **UltraSharp v2**  
+
+And (as always) many bugfixes and improvements to existing features!  
+
+[ReadMe](https://github.com/vladmandic/automatic/blob/master/README.md) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867)
+
+### Details for 2025-06-30
+
+- **Models**
+  - [Models Wiki page](https://vladmandic.github.io/sdnext-docs/Models/) is updated will all new models  
+    *note* all new image models larger than 30GB, so [offloading](https://vladmandic.github.io/sdnext-docs/Offload/) and [quantization](https://vladmandic.github.io/sdnext-docs/Quantization/) are necessary!  
+  - [OmniGen2](https://huggingface.co/OmniGen2/OmniGen2)  
+    - OmniGen2 is a powerful unified multimodal model that supports t2i and i2i workflows and uses 4B transformer with Qwen-VL-2.5 4B VLM  
+    - available via *networks -> models -> reference*  
+  - [nVidia Cosmos-Predict2 T2I](https://research.nvidia.com/labs/dir/cosmos-predict2/) *2B and 14B*  
+    - Cosmos-Predict2 T2I is a new foundational model from Nvidia in two variants: small 2B and large 14B
+    - available via *networks -> models -> reference*  
+    - *note*: 14B variant is a very large model at 36GB
+    - *note*: this is a gated model, you need to [accept terms](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image) and set your [huggingface token](https://vladmandic.github.io/sdnext-docs/Gated/)  
+  - [Black Forest Labs FLUX.1 Kontext I2I](https://bfl.ai/announcements/flux-1-kontext-dev) *Dev* variant  
+    - FLUX.1-Kontext is a 12B model billion parameter capable of editing images based on text instructions  
+    - model is primarily designed for image editing workflows, but also works for text-to-image workflows  
+    - requirements are similar to regular FLUX.1 although 2x slower  
+    - available via *networks -> models -> reference*  
+    - *note*: this is a gated model, you need to [accept terms](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) and set your [huggingface token](https://vladmandic.github.io/sdnext-docs/Gated/)  
+  - [lodestones Chroma](https://huggingface.co/lodestones/Chroma)  
+    - Chroma is a 8.9B parameter model based on *FLUX.1-schnell* and fully Apache 2.0 licensed  
+    - available via *networks -> models -> reference*  
+    - *note*: model is still in training so future updates will trigger re-download  
+    - large credits to @Trojaner for work on bringing Chroma support to SD.Next and all the optimizations around it!  
+  - [JoyCaption Beta](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava) support (in addition to existing JoyCaption Alpha)  
+    - new version of highly popular captioning model  
+    - available via *caption -> vlm caption*  
+  - [MoonDream 2](https://huggingface.co/vikhyatk/moondream2) support (updated)  
+    - really good 2B captioning model that can work on different levels of detail  
+    - available via *caption -> vlm caption*  
+  - [UltraSharp v2](https://huggingface.co/Kim2091/UltraSharpV2) support  
+    - one of the best upscalers (traditional, non-diffusion) available today!  
+    - available via *process -> upscale -> chainner*  
+- **Changes**  
+  - Update all core requirements  
+  - Support Remote VAE with *Omnigen, Lumina 2 and PixArt*  
+  - Enable quantization for captioning: *Gemma, Qwen, SMOL, Florence, JoyCaption*  
+  - Add `--trace` command line param that enables trace logging  
+  - Use Diffusers version of *OmniGen*  
+  - Control move global settings to control elements -> control settings tab  
+  - Control add setting to run hires with or without control  
+  - Update OpenVINO to 2025.2.0  
+  - Simplified and unified quantization enabled for options  
+- **SDNQ Quantization**  
+  - Add `auto` quantization mode  
+  - Add `modules_to_not_convert` support for post mode  
+  - Improve offload compatibility  
+  - Fix Qwen 2.5 with int8 matmul  
+  - Fix Dora loading  
+  - Remove per layer GC  
+  - Add support for XYZ grid to test quantization modes  
+    *note*: you need to enable quantization and choose what it applies on, then xyz grid can change quantization mode  
+    *note*: you can also enable 'add time info' to compare performance of different quantization modes  
+- **API**
+  - Add `/sdapi/v1/lora?lora=<lora_name>` endpoint that returns full lora info and metadata  
+  - Add `/sdapi/v1/controlnets?model_type=<model_type|all|None>` endpoints that returns list of available controlnets for specific model type  
+  - Set default sampler to `Default`  
+- **Fixes**  
+  - IPEX with DPM2++ FlowMatch samplers  
+  - Invalid attention processor with ControlNet  
+  - LTXVideo default scheduler  
+  - Balanced offload with OmniGen  
+  - Quantization with OmniGen  
+  - Do not save empty `params.txt` file  
+  - Override `params.txt` using `SD_PATH_PARAMS` env variable  
+  - Add `wheel` to requirements due to `pip` change  
+  - Case-insensitive sampler name matching  
+  - Fix delete file with gallery views  
+  - Add `SD_SAVE_DEBUG` env variable to report all params and metadata save operations as they happen  
+  - Fix TAESD model type detection  
+  - Fix LoRA loader incorrectly reporting errors  
+  - Fix hypertile for img2img and inpaint operations  
+  - Fix prompt parser batch size  
+  - Fix process batch with batch count  
+  - Fix process batch double image save  
+  - Fix unapply texture tiling  
+  - Fix nunchaku batch support  
+  - Fix LoRA change detection on pipeline type change  
+  - Fix LoRA load order when it includes text-encoder data  
+  - Suppress torch empty logging  
+  - Improve TAESD live preview downscale handling  
+
 ## Update for 2025-06-16
 
 - **Feature**  
   - Support for Python 3.13  
   - TeaCache support for Lumina 2  
   - Custom UNet and VAE loading support for Lumina 2  
-
 - **Changes**  
   - Increase the medvram mode threshold from 8GB to 12GB  
   - Set CPU backend to use FP32 by default  
   - Relax Python version checks for Zluda  
   - Make VAE options not require model reload  
   - Add warning about incompatible attention processors  
-
 - **Torch**  
   - Set default to `torch==2.7.1`  
   - Force upgrade pip when installing Torch  
-
 - **ROCm**  
   - Support ROCm 6.4 with `--use-nightly`  
   - Don't override user set gfx version  
   - Don't override gfx version with RX 9000  
   - Fix flash-atten repo  
-
 - **SDNQ Quantization**  
   - Add group size support for convolutional layers  
   - Add quantized matmul support for for convolutional layers  
@@ -34,7 +125,6 @@
   - Fix VAE with conv quant  
   - Don't ignore the Quantize with GPU option with offload mode `none` and `model`  
   - High VRAM usage with Lumina 2  
-
 - **Fixes**  
   - Meissonic with multiple generators  
   - OmniGen with new transformers  
diff --git a/TODO.md b/TODO.md
index 7b4b0f1f6..29a2dcf94 100644
--- a/TODO.md
+++ b/TODO.md
@@ -4,31 +4,68 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 
 ## Current
 
-### Issues/Limitations
-
-- Control: API enhance scripts compatibility  
-- Video: API support  
-
 ## Future Candidates
 
+- Refactor: Move `model_*` stuff into subfolder  
+- Refactor: sampler options  
+- Common repo for `T5` and `CLiP`  
+- Upgrade: unblock `numpy`: see `gradio`
+- Upgrade: unblock `pydantic`: see <https://github.com/Cschlaefli/automatic>
+
+### Complete Features
+
+- Python==3.13 improved support  
+- Video: API support  
+- LoRA: add OMI format support for SD35/FLUX.1  
+
+### Under Consideration
+
 - [IPAdapter negative guidance](https://github.com/huggingface/diffusers/discussions/7167)  
+- [IPAdapter composition](https://huggingface.co/ostris/ip-composition-adapter)  
+- [Refactor attention](https://github.com/huggingface/diffusers/pull/11311)  
 - [STG](https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance)  
 - [LBM](https://github.com/gojasper/LBM)  
 - [SmoothCache](https://github.com/huggingface/diffusers/issues/11135)  
-- [Magi](https://github.com/SandAI-org/MAGI-1)  
-- [SkyReels-v2](https://github.com/huggingface/diffusers/pull/11518)  
-- [WanAI-2.1 VACE](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B)  
-- [LTXVideo-0.9.7](https://github.com/huggingface/diffusers/pull/11516)  
-- [VisualClose](https://github.com/huggingface/diffusers/pull/11377)  
-- [SEVA](https://github.com/huggingface/diffusers/pull/11440)  
-- [CausVid-Plus](https://github.com/goatWu/CausVid-Plus/)  
-- [Index-AniSora](https://github.com/bilibili/Index-anisora)  
+- [MagCache](https://github.com/lllyasviel/FramePack/pull/673/files)
 - [HiDream GGUF](https://github.com/huggingface/diffusers/pull/11550)  
-- [JoyCaption-Beta-One](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava)  
 - [Diffusers guiders](https://github.com/huggingface/diffusers/pull/11311)  
 - [Nunchaku PulID](https://github.com/mit-han-lab/nunchaku/pull/274)  
-- [Dream0](https://huggingface.co/ByteDance/DreamO)  
-- [Pydantic changes](https://github.com/Cschlaefli/automatic)  
+- [Dream0 guidance](https://huggingface.co/ByteDance/DreamO)  
+- [S3Diff diffusion upscaler](https://github.com/ArcticHare105/S3Diff)  
+- [SUPIR upscaler](https://github.com/Fanghua-Yu/SUPIR)  
+
+### Monitoring
+
+- [TensorRT](https://github.com/huggingface/diffusers/pull/11173)
+
+### New models
+
+#### Stable
+- [Diffusers-0.34.0](https://github.com/huggingface/diffusers/releases/tag/v0.34.0)  
+- [WanAI-2.1 VACE](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B)(https://github.com/huggingface/diffusers/pull/11582)  
+- [LTXVideo-0.9.7](https://github.com/Lightricks/LTX-Video?tab=readme-ov-file#diffusers-integration)(https://github.com/huggingface/diffusers/pull/11516)  
+- [Cosmos-Predict2-Video](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World)(https://github.com/huggingface/diffusers/pull/11695)  
+#### Pending
+- [Magi](https://github.com/SandAI-org/MAGI-1)(https://github.com/huggingface/diffusers/pull/11713)  
+- [SEVA](https://github.com/huggingface/diffusers/pull/11440)  
+- [SkyReels-v2](https://github.com/SkyworkAI/SkyReels-V2)(https://github.com/huggingface/diffusers/pull/11518)  
+#### External:Unified/MultiModal
+- [Bagel](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT)(https://github.com/bytedance-seed/bagel)  
+- [Ming](https://github.com/inclusionAI/Ming)  
+- [Liquid](https://github.com/FoundationVision/Liquid)  
+#### External:Image2Image/Editing
+- [Step1X](https://github.com/stepfun-ai/Step1X-Edit)  
+- [SD3 UltraEdit](https://github.com/HaozheZhao/UltraEdit)  
+#### External:Video
+- [WAN2GP](https://github.com/deepbeepmeep/Wan2GP)  
+- [SelfForcing](https://github.com/guandeh17/Self-Forcing)  
+- [DiffusionForcing](https://github.com/kwsong0113/diffusion-forcing-transformer)  
+- [LanDiff](https://github.com/landiff/landiff)  
+- [HunyuanCustom](https://github.com/Tencent-Hunyuan/HunyuanCustom)  
+- [HunyuanAvatar](https://huggingface.co/tencent/HunyuanVideo-Avatar)  
+- [WAN-CausVid](https://huggingface.co/lightx2v/Wan2.1-T2V-14B-CausVid)  
+- [WAN-CausVid-Plus t2v](https://github.com/goatWu/CausVid-Plus/)  
+- [WAN-StepDistill](https://huggingface.co/lightx2v/Wan2.1-T2V-14B-StepDistill-CfgDistill)  
 
 ## Code TODO
 
@@ -47,7 +84,7 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 - model load: force-reloading entire model as loading transformers only leads to massive memory usage
 - model loader: implement model in-memory caching
 - modernui: monkey-patch for missing tabs.select event
-- modules/lora/lora_extract.py:185:9: W0511: TODO: lora: support pre-quantized flux
+- modules/lora/lora_extract.py:188:9: W0511: TODO: lora: support pre-quantized flux
 - nunchaku: batch support
 - nunchaku: cache-dir for transformer and t5 loader
 - processing: remove duplicate mask params
diff --git a/cli/model-keys.py b/cli/model-keys.py
index 45b900bd7..abab842f4 100755
--- a/cli/model-keys.py
+++ b/cli/model-keys.py
@@ -60,6 +60,8 @@ def guess_dct(dct: dict):
     if has(dct, 'model.diffusion_model.joint_blocks') and len(list(has(dct, 'model.diffusion_model.joint_blocks'))) == 38:
         return 'sd35-large'
     if has(dct, 'model.diffusion_model.double_blocks') and len(list(has(dct, 'model.diffusion_model.double_blocks'))) == 19:
+        if has(dct, 'model.diffusion_model.distilled_guidance_layer'):
+            return 'chroma'
         return 'flux-dev'
     return None
 
diff --git a/configs/chroma/model_index.json b/configs/chroma/model_index.json
new file mode 100644
index 000000000..edab1fdd6
--- /dev/null
+++ b/configs/chroma/model_index.json
@@ -0,0 +1,24 @@
+{
+  "_class_name": "ChromaPipeline",
+  "_diffusers_version": "0.34.0.dev0",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "T5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "T5Tokenizer"
+  ],
+  "transformer": [
+    "diffusers",
+    "ChromaTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}
diff --git a/configs/chroma/scheduler/scheduler_config.json b/configs/chroma/scheduler/scheduler_config.json
new file mode 100644
index 000000000..c9ef07f78
--- /dev/null
+++ b/configs/chroma/scheduler/scheduler_config.json
@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.34.0.dev0",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 3.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": true,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}
diff --git a/configs/chroma/text_encoder/config.json b/configs/chroma/text_encoder/config.json
new file mode 100644
index 000000000..483b38488
--- /dev/null
+++ b/configs/chroma/text_encoder/config.json
@@ -0,0 +1,32 @@
+{
+  "_name_or_path": "google/t5-v1_1-xxl",
+    "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": "bfloat16",
+  "transformers_version": "4.52.4",
+  "use_cache": true,
+  "vocab_size": 32128
+}
diff --git a/configs/chroma/tokenizer/added_tokens.json b/configs/chroma/tokenizer/added_tokens.json
new file mode 100644
index 000000000..3f5132007
--- /dev/null
+++ b/configs/chroma/tokenizer/added_tokens.json
@@ -0,0 +1,102 @@
+{
+  "<extra_id_0>": 32099,
+  "<extra_id_10>": 32089,
+  "<extra_id_11>": 32088,
+  "<extra_id_12>": 32087,
+  "<extra_id_13>": 32086,
+  "<extra_id_14>": 32085,
+  "<extra_id_15>": 32084,
+  "<extra_id_16>": 32083,
+  "<extra_id_17>": 32082,
+  "<extra_id_18>": 32081,
+  "<extra_id_19>": 32080,
+  "<extra_id_1>": 32098,
+  "<extra_id_20>": 32079,
+  "<extra_id_21>": 32078,
+  "<extra_id_22>": 32077,
+  "<extra_id_23>": 32076,
+  "<extra_id_24>": 32075,
+  "<extra_id_25>": 32074,
+  "<extra_id_26>": 32073,
+  "<extra_id_27>": 32072,
+  "<extra_id_28>": 32071,
+  "<extra_id_29>": 32070,
+  "<extra_id_2>": 32097,
+  "<extra_id_30>": 32069,
+  "<extra_id_31>": 32068,
+  "<extra_id_32>": 32067,
+  "<extra_id_33>": 32066,
+  "<extra_id_34>": 32065,
+  "<extra_id_35>": 32064,
+  "<extra_id_36>": 32063,
+  "<extra_id_37>": 32062,
+  "<extra_id_38>": 32061,
+  "<extra_id_39>": 32060,
+  "<extra_id_3>": 32096,
+  "<extra_id_40>": 32059,
+  "<extra_id_41>": 32058,
+  "<extra_id_42>": 32057,
+  "<extra_id_43>": 32056,
+  "<extra_id_44>": 32055,
+  "<extra_id_45>": 32054,
+  "<extra_id_46>": 32053,
+  "<extra_id_47>": 32052,
+  "<extra_id_48>": 32051,
+  "<extra_id_49>": 32050,
+  "<extra_id_4>": 32095,
+  "<extra_id_50>": 32049,
+  "<extra_id_51>": 32048,
+  "<extra_id_52>": 32047,
+  "<extra_id_53>": 32046,
+  "<extra_id_54>": 32045,
+  "<extra_id_55>": 32044,
+  "<extra_id_56>": 32043,
+  "<extra_id_57>": 32042,
+  "<extra_id_58>": 32041,
+  "<extra_id_59>": 32040,
+  "<extra_id_5>": 32094,
+  "<extra_id_60>": 32039,
+  "<extra_id_61>": 32038,
+  "<extra_id_62>": 32037,
+  "<extra_id_63>": 32036,
+  "<extra_id_64>": 32035,
+  "<extra_id_65>": 32034,
+  "<extra_id_66>": 32033,
+  "<extra_id_67>": 32032,
+  "<extra_id_68>": 32031,
+  "<extra_id_69>": 32030,
+  "<extra_id_6>": 32093,
+  "<extra_id_70>": 32029,
+  "<extra_id_71>": 32028,
+  "<extra_id_72>": 32027,
+  "<extra_id_73>": 32026,
+  "<extra_id_74>": 32025,
+  "<extra_id_75>": 32024,
+  "<extra_id_76>": 32023,
+  "<extra_id_77>": 32022,
+  "<extra_id_78>": 32021,
+  "<extra_id_79>": 32020,
+  "<extra_id_7>": 32092,
+  "<extra_id_80>": 32019,
+  "<extra_id_81>": 32018,
+  "<extra_id_82>": 32017,
+  "<extra_id_83>": 32016,
+  "<extra_id_84>": 32015,
+  "<extra_id_85>": 32014,
+  "<extra_id_86>": 32013,
+  "<extra_id_87>": 32012,
+  "<extra_id_88>": 32011,
+  "<extra_id_89>": 32010,
+  "<extra_id_8>": 32091,
+  "<extra_id_90>": 32009,
+  "<extra_id_91>": 32008,
+  "<extra_id_92>": 32007,
+  "<extra_id_93>": 32006,
+  "<extra_id_94>": 32005,
+  "<extra_id_95>": 32004,
+  "<extra_id_96>": 32003,
+  "<extra_id_97>": 32002,
+  "<extra_id_98>": 32001,
+  "<extra_id_99>": 32000,
+  "<extra_id_9>": 32090
+}
diff --git a/configs/chroma/tokenizer/special_tokens_map.json b/configs/chroma/tokenizer/special_tokens_map.json
new file mode 100644
index 000000000..17ade346a
--- /dev/null
+++ b/configs/chroma/tokenizer/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
+  }
+}
diff --git a/configs/chroma/tokenizer/spiece.model b/configs/chroma/tokenizer/spiece.model
new file mode 100644
index 000000000..4e28ff6eb
Binary files /dev/null and b/configs/chroma/tokenizer/spiece.model differ
diff --git a/configs/chroma/tokenizer/tokenizer_config.json b/configs/chroma/tokenizer/tokenizer_config.json
new file mode 100644
index 000000000..b336fa236
--- /dev/null
+++ b/configs/chroma/tokenizer/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": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<extra_id_98>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<extra_id_97>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<extra_id_96>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<extra_id_95>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<extra_id_94>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<extra_id_93>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<extra_id_92>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<extra_id_91>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32009": {
+      "content": "<extra_id_90>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32010": {
+      "content": "<extra_id_89>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32011": {
+      "content": "<extra_id_88>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32012": {
+      "content": "<extra_id_87>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32013": {
+      "content": "<extra_id_86>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32014": {
+      "content": "<extra_id_85>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32015": {
+      "content": "<extra_id_84>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32016": {
+      "content": "<extra_id_83>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32017": {
+      "content": "<extra_id_82>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32018": {
+      "content": "<extra_id_81>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32019": {
+      "content": "<extra_id_80>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32020": {
+      "content": "<extra_id_79>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32021": {
+      "content": "<extra_id_78>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32022": {
+      "content": "<extra_id_77>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32023": {
+      "content": "<extra_id_76>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32024": {
+      "content": "<extra_id_75>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32025": {
+      "content": "<extra_id_74>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32026": {
+      "content": "<extra_id_73>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32027": {
+      "content": "<extra_id_72>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32028": {
+      "content": "<extra_id_71>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32029": {
+      "content": "<extra_id_70>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32030": {
+      "content": "<extra_id_69>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32031": {
+      "content": "<extra_id_68>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32032": {
+      "content": "<extra_id_67>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32033": {
+      "content": "<extra_id_66>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32034": {
+      "content": "<extra_id_65>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32035": {
+      "content": "<extra_id_64>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32036": {
+      "content": "<extra_id_63>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32037": {
+      "content": "<extra_id_62>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32038": {
+      "content": "<extra_id_61>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32039": {
+      "content": "<extra_id_60>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32040": {
+      "content": "<extra_id_59>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32041": {
+      "content": "<extra_id_58>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32042": {
+      "content": "<extra_id_57>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32043": {
+      "content": "<extra_id_56>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32044": {
+      "content": "<extra_id_55>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32045": {
+      "content": "<extra_id_54>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32046": {
+      "content": "<extra_id_53>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32047": {
+      "content": "<extra_id_52>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32048": {
+      "content": "<extra_id_51>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32049": {
+      "content": "<extra_id_50>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32050": {
+      "content": "<extra_id_49>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32051": {
+      "content": "<extra_id_48>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32052": {
+      "content": "<extra_id_47>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32053": {
+      "content": "<extra_id_46>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32054": {
+      "content": "<extra_id_45>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32055": {
+      "content": "<extra_id_44>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32056": {
+      "content": "<extra_id_43>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32057": {
+      "content": "<extra_id_42>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32058": {
+      "content": "<extra_id_41>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32059": {
+      "content": "<extra_id_40>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32060": {
+      "content": "<extra_id_39>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32061": {
+      "content": "<extra_id_38>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32062": {
+      "content": "<extra_id_37>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32063": {
+      "content": "<extra_id_36>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32064": {
+      "content": "<extra_id_35>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32065": {
+      "content": "<extra_id_34>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32066": {
+      "content": "<extra_id_33>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32067": {
+      "content": "<extra_id_32>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32068": {
+      "content": "<extra_id_31>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32069": {
+      "content": "<extra_id_30>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32070": {
+      "content": "<extra_id_29>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32071": {
+      "content": "<extra_id_28>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32072": {
+      "content": "<extra_id_27>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32073": {
+      "content": "<extra_id_26>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32074": {
+      "content": "<extra_id_25>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32075": {
+      "content": "<extra_id_24>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32076": {
+      "content": "<extra_id_23>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32077": {
+      "content": "<extra_id_22>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32078": {
+      "content": "<extra_id_21>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32079": {
+      "content": "<extra_id_20>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32080": {
+      "content": "<extra_id_19>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32081": {
+      "content": "<extra_id_18>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32082": {
+      "content": "<extra_id_17>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32083": {
+      "content": "<extra_id_16>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32084": {
+      "content": "<extra_id_15>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32085": {
+      "content": "<extra_id_14>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32086": {
+      "content": "<extra_id_13>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32087": {
+      "content": "<extra_id_12>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32088": {
+      "content": "<extra_id_11>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32089": {
+      "content": "<extra_id_10>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32090": {
+      "content": "<extra_id_9>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32091": {
+      "content": "<extra_id_8>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32092": {
+      "content": "<extra_id_7>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32093": {
+      "content": "<extra_id_6>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32094": {
+      "content": "<extra_id_5>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32095": {
+      "content": "<extra_id_4>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32096": {
+      "content": "<extra_id_3>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32097": {
+      "content": "<extra_id_2>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32098": {
+      "content": "<extra_id_1>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32099": {
+      "content": "<extra_id_0>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "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>"
+}
diff --git a/configs/chroma/transformer/config.json b/configs/chroma/transformer/config.json
new file mode 100644
index 000000000..86d6daa85
--- /dev/null
+++ b/configs/chroma/transformer/config.json
@@ -0,0 +1,20 @@
+{
+  "_class_name": "ChromaTransformer2DModel",
+  "_diffusers_version": "0.34.0.dev0",
+  "approximator_hidden_dim": 5120,
+  "approximator_in_factor": 16,
+  "approximator_layers": 5,
+  "attention_head_dim": 128,
+  "axes_dims_rope": [
+    16,
+    56,
+    56
+  ],
+  "in_channels": 64,
+  "joint_attention_dim": 4096,
+  "num_attention_heads": 24,
+  "num_layers": 19,
+  "num_single_layers": 38,
+  "out_channels": null,
+  "patch_size": 1
+}
diff --git a/configs/chroma/transformer/diffusion_pytorch_model.safetensors.index.json b/configs/chroma/transformer/diffusion_pytorch_model.safetensors.index.json
new file mode 100644
index 000000000..8a52f5e7b
--- /dev/null
+++ b/configs/chroma/transformer/diffusion_pytorch_model.safetensors.index.json
@@ -0,0 +1,1030 @@
+{
+  "metadata": {
+    "total_size": 17799966848
+  },
+  "weight_map": {
+    "context_embedder.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "context_embedder.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.in_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.in_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.0.linear_1.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.0.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.0.linear_2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.0.linear_2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.1.linear_1.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.1.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.1.linear_2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.1.linear_2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.2.linear_1.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.2.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.2.linear_2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.2.linear_2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.3.linear_1.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.3.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.3.linear_2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.3.linear_2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.4.linear_1.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.4.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.4.linear_2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.layers.4.linear_2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.norms.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.norms.1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.norms.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.norms.3.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.norms.4.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.out_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "distilled_guidance_layer.out_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.proj_mlp.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.proj_mlp.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.proj_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.0.proj_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.proj_mlp.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.proj_mlp.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.proj_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.1.proj_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.10.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.11.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.12.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.13.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.14.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.15.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.16.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.17.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.18.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.19.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.proj_mlp.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.proj_mlp.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.proj_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.2.proj_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.20.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.21.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.22.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.23.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.24.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.25.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.26.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.27.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.28.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.29.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.proj_mlp.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.3.proj_mlp.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "single_transformer_blocks.3.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.3.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.30.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.31.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.32.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.33.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.34.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.35.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.36.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.37.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.4.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.5.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.6.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.7.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.8.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.proj_mlp.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.proj_mlp.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "single_transformer_blocks.9.proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "transformer_blocks.0.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.0.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.1.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.10.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.11.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.12.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.13.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.14.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.15.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.16.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.17.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.18.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.2.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.3.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.4.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.5.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.6.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.7.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.8.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.add_k_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.add_q_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.add_v_proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_add_out.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_k.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_out.0.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_q.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_v.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff_context.net.2.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "transformer_blocks.9.ff_context.net.2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "x_embedder.bias": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "x_embedder.weight": "diffusion_pytorch_model-00001-of-00002.safetensors"
+  }
+}
diff --git a/configs/chroma/vae/config.json b/configs/chroma/vae/config.json
new file mode 100644
index 000000000..188c62140
--- /dev/null
+++ b/configs/chroma/vae/config.json
@@ -0,0 +1,37 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.34.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,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 0.3611,
+  "shift_factor": 0.1159,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": false,
+  "use_quant_conv": false
+}
diff --git a/extensions-builtin/Lora/lora_extract.py b/extensions-builtin/Lora/lora_extract.py
index 85bcc2c95..c6b65fca7 100644
--- a/extensions-builtin/Lora/lora_extract.py
+++ b/extensions-builtin/Lora/lora_extract.py
@@ -107,6 +107,9 @@ def make_meta(fn, maxrank, rank_ratio):
     elif shared.sd_model_type == "f1":
         meta["model_spec.architecture"] = "flux-1-dev/lora"
         meta["ss_base_model_version"] = "flux1"
+    elif shared.sd_model_type == "chroma":
+        meta["model_spec.architecture"] = "chroma/lora"
+        meta["ss_base_model_version"] = "chroma"
     elif shared.sd_model_type == "sc":
         meta["model_spec.architecture"] = "stable-cascade-v1-prior/lora"
     return meta
diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index e59555993..0a4657383 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -139,7 +139,7 @@ def load_network(name, network_on_disk) -> network.Network:
     net = network.Network(name, network_on_disk)
     net.mtime = os.path.getmtime(network_on_disk.filename)
     sd = sd_models.read_state_dict(network_on_disk.filename, what='network')
-    if shared.sd_model_type == 'f1':  # if kohya flux lora, convert state_dict
+    if shared.sd_model_type in ['f1', 'chroma']:  # if kohya flux lora, convert state_dict
         sd = lora_convert._convert_kohya_flux_lora_to_diffusers(sd) or sd  # pylint: disable=protected-access
     assign_network_names_to_compvis_modules(shared.sd_model)
     keys_failed_to_match = {}
diff --git a/extensions-builtin/sd-extension-chainner b/extensions-builtin/sd-extension-chainner
index a841a9dfe..c12e8cda4 160000
--- a/extensions-builtin/sd-extension-chainner
+++ b/extensions-builtin/sd-extension-chainner
@@ -1 +1 @@
-Subproject commit a841a9dfeab8178d0e0236e156a11973cca9462d
+Subproject commit c12e8cda4e45a1d5f20659a30e329c7cc7e69339
diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui
index ab51d26c8..6ad9a291b 160000
--- a/extensions-builtin/sdnext-modernui
+++ b/extensions-builtin/sdnext-modernui
@@ -1 +1 @@
-Subproject commit ab51d26c83d051ab6af2bf8647b36507e626787e
+Subproject commit 6ad9a291b515041adbefdb3eb94f87fb2892c902
diff --git a/html/reference.json b/html/reference.json
index 2441112e6..e90aad1d1 100644
--- a/html/reference.json
+++ b/html/reference.json
@@ -158,24 +158,18 @@
     "skip": true,
     "extras": "sampler: Default, cfg_scale: 3.5"
   },
-  "Black Forest Labs FLUX.1 Dev qint8": {
-    "path": "Disty0/FLUX.1-dev-qint8",
-    "preview": "black-forest-labs--FLUX.1-dev.jpg",
-    "desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching",
+  "Black Forest Labs FLUX.1 Kontext Dev": {
+    "path": "black-forest-labs/FLUX.1-Kontext-dev",
+    "preview": "black-forest-labs--FLUX.1-Kontext-dev.jpg",
+    "desc": "FLUX.1 Kontext [dev] is a 12 billion parameter rectified flow transformer capable of editing images based on text instructions.",
     "skip": true,
     "extras": "sampler: Default, cfg_scale: 3.5"
   },
-  "Black Forest Labs FLUX.1 Dev qint4": {
-    "path": "Disty0/FLUX.1-dev-qint4",
-    "preview": "black-forest-labs--FLUX.1-dev.jpg",
-    "desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching",
-    "skip": true,
-    "extras": "sampler: Default, cfg_scale: 3.5"
-  },
-  "Black Forest Labs FLUX.1 Dev nf4": {
-    "path": "sayakpaul/flux.1-dev-nf4",
-    "preview": "black-forest-labs--FLUX.1-dev.jpg",
-    "desc": "FLUX.1 models are based on a hybrid architecture of multimodal and parallel diffusion transformer blocks, scaled to 12B parameters and builing on flow matching",
+
+  "lodestones Chroma": {
+    "path": "lodestones/Chroma",
+    "preview": "lodestones--Chroma.jpg",
+    "desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. The model is still training right now, and I’d love to hear your thoughts! Your input and feedback are really appreciated.",
     "skip": true,
     "extras": "sampler: Default, cfg_scale: 3.5"
   },
@@ -236,15 +230,35 @@
     "desc": "Sana is a text-to-image framework that can efficiently generate images up to 4096 × 4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU.",
     "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
     "skip": true
-  }, 
+  },
+
+  "nVidia Cosmos-Predict2 T2I 2B": {
+    "path": "nvidia/Cosmos-Predict2-2B-Text2Image",
+    "desc": "Cosmos-Predict2: A family of highly performant pre-trained world foundation models purpose-built for generating physics-aware images, videos and world states for physical AI development.",
+    "preview": "nvidia--Cosmos-Predict2-2B-Text2Image.jpg",
+    "skip": true
+  },
+  "nVidia Cosmos-Predict2 T2I 14B": {
+    "path": "nvidia/Cosmos-Predict2-14B-Text2Image",
+    "desc": "Cosmos-Predict2: A family of highly performant pre-trained world foundation models purpose-built for generating physics-aware images, videos and world states for physical AI development.",
+    "preview": "nvidia--Cosmos-Predict2-2B-Text2Image.jpg",
+    "skip": true
+  },
 
   "VectorSpaceLab OmniGen v1": {
-    "path": "Shitao/OmniGen-v1",
+    "path": "Shitao/OmniGen-v1-diffusers",
     "desc": "OmniGen is a unified image generation model that can generate a wide range of images from multi-modal prompts. It is designed to be simple, flexible and easy to use.",
     "preview": "Shitao--OmniGen-v1.jpg",
     "skip": true
   }, 
 
+  "VectorSpaceLab OmniGen v2": {
+    "path": "OmniGen2/OmniGen2",
+    "desc": "OmniGen2 is a powerful and efficient unified multimodal model. Unlike OmniGen v1, OmniGen2 features two distinct decoding pathways for text and image modalities, utilizing unshared parameters and a decoupled image tokenizer.",
+    "preview": "OmniGen2--OmniGen2.jpg",
+    "skip": true
+  }, 
+
   "AuraFlow 0.3": {
     "path": "fal/AuraFlow-v0.3",
     "desc": "AuraFlow v0.3 is the fully open-sourced flow-based text-to-image generation model. The model was trained with more compute compared to the previous version, AuraFlow-v0.2. Compared to AuraFlow-v0.2, the model is fine-tuned on more aesthetic datasets and now supports various aspect ratio, (now width and height up to 1536 pixels).",
diff --git a/installer.py b/installer.py
index bf4d0b408..4de6f8103 100644
--- a/installer.py
+++ b/installer.py
@@ -129,6 +129,13 @@ def setup_logging():
         def get(self):
             return self.buffer
 
+    class LogFilter(logging.Filter):
+        def __init__(self):
+            super().__init__()
+
+        def filter(self, record):
+            return len(record.getMessage()) > 2
+
     t_start = time.time()
     from functools import partial, partialmethod
     from logging.handlers import RotatingFileHandler
@@ -147,7 +154,7 @@ def setup_logging():
     logging.Logger.trace = partialmethod(logging.Logger.log, logging.TRACE)
     logging.trace = partial(logging.log, logging.TRACE)
 
-    level = logging.DEBUG if args.debug else logging.INFO
+    level = logging.DEBUG if (args.debug or args.trace) else logging.INFO
     log.setLevel(logging.DEBUG) # log to file is always at level debug for facility `sd`
     log.print = rprint
     global console # pylint: disable=global-statement
@@ -173,22 +180,30 @@ def setup_logging():
     while log.hasHandlers() and len(log.handlers) > 0:
         log.removeHandler(log.handlers[0])
 
+    log_filter = LogFilter()
     # handlers
     rh = RichHandler(show_time=True, omit_repeated_times=False, show_level=True, show_path=False, markup=False, rich_tracebacks=True, log_time_format='%H:%M:%S-%f', level=level, console=console)
+    if args.trace:
+        rh.formatter = logging.Formatter('[%(module)s][%(pathname)s:%(lineno)d]  %(message)s')
+    rh.addFilter(log_filter)
     rh.setLevel(level)
     log.addHandler(rh)
 
     fh = RotatingFileHandler(log_file, maxBytes=32*1024*1024, backupCount=9, encoding='utf-8', delay=True) # 10MB default for log rotation
+    if args.trace:
+        fh.formatter = logging.Formatter(f'%(asctime)s | {hostname} | %(name)s | %(levelname)s | %(module)s | | %(pathname)s:%(lineno)d | %(message)s')
+    else:
+        fh.formatter = logging.Formatter(f'%(asctime)s | {hostname} | %(name)s | %(levelname)s | %(module)s | %(message)s')
+    fh.addFilter(log_filter)
+    fh.setLevel(logging.DEBUG)
+    log.addHandler(fh)
     global log_rolled # pylint: disable=global-statement
     if not log_rolled and args.debug and not args.log:
         fh.doRollover()
         log_rolled = True
 
-    fh.formatter = logging.Formatter(f'%(asctime)s | {hostname} | %(name)s | %(levelname)s | %(module)s | %(message)s')
-    fh.setLevel(logging.DEBUG)
-    log.addHandler(fh)
-
     rb = RingBuffer(100) # 100 entries default in log ring buffer
+    rb.addFilter(log_filter)
     rb.setLevel(level)
     log.addHandler(rb)
     log.buffer = rb.buffer
@@ -520,9 +535,9 @@ def check_python(supported_minors=[], experimental_minors=[], reason=None):
     log.info(f'Python: version={platform.python_version()} platform={platform.system()} bin="{sys.executable}" venv="{sys.prefix}"')
     if not (int(sys.version_info.major) == 3 and int(sys.version_info.minor) in supported_minors):
         if (int(sys.version_info.major) == 3 and int(sys.version_info.minor) in experimental_minors):
-            log.warning(f"Python version experimental: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} recommended 3.{supported_minors}")
+            log.warning(f"Python experimental: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}")
         else:
-            log.error(f"Python version incompatible: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}")
+            log.error(f"Python incompatible: current {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}")
             if reason is not None:
                 log.error(reason)
             if not args.ignore and not args.experimental:
@@ -546,7 +561,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git or args.experimental:
         return
-    sha = '8adc6003ba4dbf5b61bb4f1ce571e9e55e145a99' # diffusers commit hash
+    sha = '00f95b9755718aabb65456e791b8408526ae6e76' # diffusers commit hash
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
     cur = opts.get('diffusers_version', '') if minor > 0 else ''
@@ -752,8 +767,8 @@ def install_openvino():
     else:
         torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+cpu torchvision==0.22.1+cpu --index-url https://download.pytorch.org/whl/cpu')
 
-    install(os.environ.get('OPENVINO_COMMAND', 'openvino==2025.1.0'), 'openvino')
-    install(os.environ.get('NNCF_COMMAND', 'nncf==2.16.0'), 'nncf')
+    install(os.environ.get('OPENVINO_COMMAND', 'openvino==2025.2.0'), 'openvino')
+    install(os.environ.get('NNCF_COMMAND', 'nncf==2.17.0'), 'nncf')
     os.environ.setdefault('PYTORCH_TRACING_MODE', 'TORCHFX')
     if os.environ.get("NEOReadDebugKeys", None) is None:
         os.environ.setdefault('NEOReadDebugKeys', '1')
@@ -1528,7 +1543,8 @@ def add_args(parser):
 
     group_log = parser.add_argument_group('Logging')
     group_log.add_argument("--log", type=str, default=os.environ.get("SD_LOG", None), help="Set log file, default: %(default)s")
-    group_log.add_argument('--debug', default=os.environ.get("SD_DEBUG",False), action='store_true', help="Run installer with debug logging, default: %(default)s")
+    group_log.add_argument('--debug', default=os.environ.get("SD_DEBUG",False), action='store_true', help="Run with debug logging, default: %(default)s")
+    group_log.add_argument("--trace", default=os.environ.get("SD_TRACE", False), action='store_true', help="Run with trace logging, default: %(default)s")
     group_log.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s")
     group_log.add_argument('--docs', default=os.environ.get("SD_DOCS", False), action='store_true', help="Mount API docs, default: %(default)s")
     group_log.add_argument("--api-log", default=os.environ.get("SD_APILOG", False), action='store_true', help="Log all API requests")
diff --git a/javascript/gallery.js b/javascript/gallery.js
index 32ae0bde4..eaa646c66 100644
--- a/javascript/gallery.js
+++ b/javascript/gallery.js
@@ -131,7 +131,7 @@ class GalleryFile extends HTMLElement {
     }
     const ext = this.name.split('.').pop().toLowerCase();
     if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'jxl', 'svg', 'mp4'].includes(ext)) {
-      console.error(`gallery: type=${ext} file=${this.name} unsupported`);
+      // console.error(`gallery: type=${ext} file=${this.name} unsupported`);
       return;
     }
     this.hash = await getHash(`${this.folder}/${this.name}/${this.size}/${this.mtime}`); // eslint-disable-line no-use-before-define
@@ -286,7 +286,7 @@ async function gallerySearch(evt) {
 
 const findDuplicates = (arr, key) => {
   const map = new Map();
-  return arr.filter(item => {
+  return arr.filter((item) => {
     const value = item[key];
     if (map.has(value)) return true;
     map.set(value, true);
diff --git a/javascript/ui.js b/javascript/ui.js
index b8112d523..ea8af5902 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -60,6 +60,17 @@ function selected_gallery_index() {
   return result;
 }
 
+function selected_gallery_files() {
+  let allImages = [];
+  try {
+    let allCurrentButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnail-item.thumbnail-small');
+    if (allCurrentButtons.length === 0) allCurrentButtons = gradioApp().querySelectorAll('.gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small');
+    allImages = Array.from(allCurrentButtons).map((v) => v.querySelector('img')?.src);
+  } catch { /**/ }
+  const selectedIndex = selected_gallery_index();
+  return [allImages, selectedIndex];
+}
+
 function extract_image_from_gallery(gallery) {
   if (gallery.length === 0) return [null];
   if (gallery.length === 1) return [gallery[0]];
diff --git a/models/Reference/OmniGen2--OmniGen2.jpg b/models/Reference/OmniGen2--OmniGen2.jpg
new file mode 100644
index 000000000..d8fce4558
Binary files /dev/null and b/models/Reference/OmniGen2--OmniGen2.jpg differ
diff --git a/models/Reference/black-forest-labs--FLUX.1-Kontext-dev.jpg b/models/Reference/black-forest-labs--FLUX.1-Kontext-dev.jpg
new file mode 100644
index 000000000..04521fdd2
Binary files /dev/null and b/models/Reference/black-forest-labs--FLUX.1-Kontext-dev.jpg differ
diff --git a/models/Reference/lodestones--Chroma.jpg b/models/Reference/lodestones--Chroma.jpg
index 6c0c5d2f4..78e6e33f7 100644
Binary files a/models/Reference/lodestones--Chroma.jpg and b/models/Reference/lodestones--Chroma.jpg differ
diff --git a/models/Reference/nvidia--Cosmos-Predict2-2B-Text2Image.jpg b/models/Reference/nvidia--Cosmos-Predict2-2B-Text2Image.jpg
new file mode 100644
index 000000000..54b9af8c1
Binary files /dev/null and b/models/Reference/nvidia--Cosmos-Predict2-2B-Text2Image.jpg differ
diff --git a/modules/api/api.py b/modules/api/api.py
index 72a2090a0..ad26a384a 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -5,7 +5,7 @@ from fastapi import FastAPI, APIRouter, Depends, Request
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from fastapi.exceptions import HTTPException
 from modules import errors, shared, postprocessing
-from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery, docs
+from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery, loras, docs
 
 
 errors.install()
@@ -78,6 +78,7 @@ class Api:
         self.add_api_route("/sdapi/v1/samplers", endpoints.get_samplers, methods=["GET"], response_model=List[models.ItemSampler])
         self.add_api_route("/sdapi/v1/upscalers", endpoints.get_upscalers, methods=["GET"], response_model=List[models.ItemUpscaler])
         self.add_api_route("/sdapi/v1/sd-models", endpoints.get_sd_models, methods=["GET"], response_model=List[models.ItemModel])
+        self.add_api_route("/sdapi/v1/controlnets", endpoints.get_controlnets, methods=["GET"], response_model=List[str])
         self.add_api_route("/sdapi/v1/hypernetworks", endpoints.get_hypernetworks, methods=["GET"], response_model=List[models.ItemHypernetwork])
         self.add_api_route("/sdapi/v1/face-restorers", endpoints.get_detailers, methods=["GET"], response_model=List[models.ItemDetailer])
         self.add_api_route("/sdapi/v1/prompt-styles", endpoints.get_prompt_styles, methods=["GET"], response_model=List[models.ItemStyle])
@@ -100,8 +101,9 @@ class Api:
 
         # lora api
         if shared.native:
-            self.add_api_route("/sdapi/v1/loras", endpoints.get_loras, methods=["GET"], response_model=List[dict])
-            self.add_api_route("/sdapi/v1/refresh-loras", endpoints.post_refresh_loras, methods=["POST"])
+            self.add_api_route("/sdapi/v1/lora", loras.get_lora, methods=["GET"], response_model=dict)
+            self.add_api_route("/sdapi/v1/loras", loras.get_loras, methods=["GET"], response_model=List[dict])
+            self.add_api_route("/sdapi/v1/refresh-loras", loras.post_refresh_loras, methods=["POST"])
 
         # gallery api
         gallery.register_api(self.app)
diff --git a/modules/api/control.py b/modules/api/control.py
index 9d6512cd2..9e06d874a 100644
--- a/modules/api/control.py
+++ b/modules/api/control.py
@@ -22,7 +22,7 @@ ReqControl = models.create_model_from_signature(
     func = run.control_run,
     model_name = "StableDiffusionProcessingControl",
     additional_fields = [
-        {"key": "sampler_name", "type": str, "default": "UniPC"},
+        {"key": "sampler_name", "type": str, "default": "Default"},
         {"key": "script_name", "type": Optional[str], "default": None},
         {"key": "script_args", "type": list, "default": []},
         {"key": "send_images", "type": bool, "default": True},
diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py
index 80b46f324..561afafe0 100644
--- a/modules/api/endpoints.py
+++ b/modules/api/endpoints.py
@@ -23,6 +23,10 @@ def get_sd_models():
         checkpoints.append({"title": v.title, "model_name": v.name, "filename": v.filename, "type": v.type, "hash": v.shorthash, "sha256": v.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(v)})
     return checkpoints
 
+def get_controlnets(model_type: Optional[str] = None):
+    from modules.control.units.controlnet import api_list_models
+    return api_list_models(model_type)
+
 def get_hypernetworks():
     return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
 
@@ -43,12 +47,6 @@ def get_embeddings():
 
     return {"loaded": convert_embeddings(db.word_embeddings), "skipped": convert_embeddings(db.skipped_embeddings)}
 
-def get_loras():
-    from modules.lora import network, lora_load
-    def create_lora_json(obj: network.NetworkOnDisk):
-        return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata }
-    return [create_lora_json(obj) for obj in lora_load.available_networks.values()]
-
 def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, filename: Optional[str] = None, title: Optional[str] = None, fullname: Optional[str] = None, hash: Optional[str] = None): # pylint: disable=redefined-builtin
     res = []
     for pg in shared.extra_networks:
@@ -158,10 +156,6 @@ def post_refresh_vae():
     shared.refresh_vaes()
     return {}
 
-def post_refresh_loras():
-    from modules.lora import lora_load
-    return lora_load.list_available_networks()
-
 def get_extensions_list():
     from modules import extensions
     extensions.list_extensions()
diff --git a/modules/api/gallery.py b/modules/api/gallery.py
index e1add81af..6510cd673 100644
--- a/modules/api/gallery.py
+++ b/modules/api/gallery.py
@@ -6,7 +6,7 @@ from typing import List, Union
 from urllib.parse import quote, unquote
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
-from starlette.websockets import WebSocket, WebSocketState, WebSocketDisconnect
+from starlette.websockets import WebSocket, WebSocketState
 from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
 from PIL import Image
 from modules import shared, images, files_cache
@@ -196,6 +196,6 @@ def register_api(app: FastAPI): # register api
             await manager.send(ws, '#END#')
             t1 = time.time()
             shared.log.debug(f'Gallery: type=ws folder="{folder}" files={numFiles} time={t1-t0:.3f}')
-        except WebSocketDisconnect:
-            debug('Browser WS unexpected disconnect')
+        except Exception as e:
+            debug(f'Browser WS error: {e}')
         manager.disconnect(ws)
diff --git a/modules/api/loras.py b/modules/api/loras.py
new file mode 100644
index 000000000..7e65a709d
--- /dev/null
+++ b/modules/api/loras.py
@@ -0,0 +1,21 @@
+from fastapi.exceptions import HTTPException
+
+
+def get_lora(lora: str) -> dict:
+    from modules.lora import lora_load
+    if lora not in lora_load.available_networks:
+        raise HTTPException(status_code=404, detail=f"Lora '{lora}' not found")
+    obj = lora_load.available_networks[lora]
+    obj.info = obj.get_info()
+    obj.desc = obj.get_desc()
+    return obj.__dict__
+
+def get_loras():
+    from modules.lora import network, lora_load
+    def create_lora_json(obj: network.NetworkOnDisk):
+        return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata }
+    return [create_lora_json(obj) for obj in lora_load.available_networks.values()]
+
+def post_refresh_loras():
+    from modules.lora import lora_load
+    return lora_load.list_available_networks()
diff --git a/modules/api/models.py b/modules/api/models.py
index f463a2dd6..a19c59858 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -198,7 +198,7 @@ ReqTxt2Img = PydanticModelGenerator(
     StableDiffusionProcessingTxt2Img,
     [
         {"key": "sampler_index", "type": Union[int, str], "default": 0},
-        {"key": "sampler_name", "type": str, "default": "UniPC"},
+        {"key": "sampler_name", "type": str, "default": "Default"},
         {"key": "hr_sampler_name", "type": str, "default": "Same as primary"},
         {"key": "script_name", "type": Optional[str], "default": "none"},
         {"key": "script_args", "type": list, "default": []},
diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py
index d2e59348c..5e9e372af 100644
--- a/modules/control/units/controlnet.py
+++ b/modules/control/units/controlnet.py
@@ -137,6 +137,23 @@ def find_models():
 
 find_models()
 
+
+def api_list_models(model_type: str = None):
+    import modules.shared
+    model_type = model_type or modules.shared.sd_model_type
+    model_list = []
+    if model_type == 'sd' or model_type == 'all':
+        model_list += list(predefined_sd15)
+    if model_type == 'sdxl' or model_type == 'all':
+        model_list += list(predefined_sdxl)
+    if model_type == 'f1' or model_type == 'all':
+        model_list += list(predefined_f1)
+    if model_type == 'sd3' or model_type == 'all':
+        model_list += list(predefined_sd3)
+    model_list += sorted(find_models())
+    return model_list
+
+
 def list_models(refresh=False):
     import modules.shared
     global models # pylint: disable=global-statement
@@ -286,21 +303,29 @@ class ControlNet():
                     return
                 if self.dtype is not None:
                     self.model.to(self.dtype)
-                if "ControlNet" in opts.sdnq_quantize_weights:
+                if "Control" in opts.sdnq_quantize_weights:
                     try:
                         log.debug(f'Control {what} model SDNQ Compress: id="{model_id}"')
                         from modules.model_quant import sdnq_quantize_model
                         self.model = sdnq_quantize_model(self.model)
                     except Exception as e:
                         log.error(f'Control {what} model SDNQ Compression failed: id="{model_id}" {e}')
-                elif "ControlNet" in opts.optimum_quanto_weights:
+                elif "Control" in opts.optimum_quanto_weights:
                     try:
                         log.debug(f'Control {what} model Optimum Quanto: id="{model_id}"')
-                        model_quant.load_quanto('Load model: type=ControlNet')
+                        model_quant.load_quanto('Load model: type=Control')
                         from modules.model_quant import optimum_quanto_model
                         self.model = optimum_quanto_model(self.model)
                     except Exception as e:
                         log.error(f'Control {what} model Optimum Quanto: id="{model_id}" {e}')
+                elif "Control" in opts.torchao_quantization:
+                    try:
+                        log.debug(f'Control {what} model Torch AO: id="{model_id}"')
+                        model_quant.load_torchao('Load model: type=Control')
+                        from modules.model_quant import torchao_quantization
+                        self.model = torchao_quantization(self.model)
+                    except Exception as e:
+                        log.error(f'Control {what} model Torch AO: id="{model_id}" {e}')
                 if self.device is not None:
                     self.model.to(self.device)
                 t1 = time.time()
diff --git a/modules/devices.py b/modules/devices.py
index 1c35f2683..8d15fd238 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -665,6 +665,6 @@ def normalize_device(dev):
 
 
 def same_device(d1, d2):
-    if d1.type != d2.type:
+    if torch.device(d1).type != torch.device(d2).type:
         return False
     return normalize_device(d1) == normalize_device(d2)
diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index 98c2b39cd..6393c9051 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -3,7 +3,7 @@ import io
 import os
 from PIL import Image
 import gradio as gr
-from modules.paths import data_path
+from modules.paths import params_path
 from modules import shared, gr_tempdir, script_callbacks, images
 from modules.infotext import parse, mapping, quote, unquote # pylint: disable=unused-import
 
@@ -223,9 +223,8 @@ def connect_paste(button, local_paste_fields, input_comp, override_settings_comp
 
     def paste_func(prompt):
         if prompt is None or len(prompt.strip()) == 0:
-            filename = os.path.join(data_path, "params.txt")
-            if os.path.exists(filename):
-                with open(filename, "r", encoding="utf8") as file:
+            if os.path.exists(params_path):
+                with open(params_path, "r", encoding="utf8") as file:
                     prompt = file.read()
                 shared.log.debug(f'Prompt parse: type="params" prompt="{prompt}"')
             else:
diff --git a/modules/gr_tempdir.py b/modules/gr_tempdir.py
index bbe2b2192..f19fd53a3 100644
--- a/modules/gr_tempdir.py
+++ b/modules/gr_tempdir.py
@@ -74,8 +74,9 @@ def pil_to_temp_file(self, img: Image, dir: str, format="png") -> str: # pylint:
         shared.state.image_history += 1
     params = ', '.join([f'{k}: {v}' for k, v in img.info.items()])
     params = params[12:] if params.startswith('parameters: ') else params
-    with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
-        file.write(params)
+    if len(params) > 2:
+        with open(paths.params_path, "w", encoding="utf8") as file:
+            file.write(params)
     return name
 
 
diff --git a/modules/hidiffusion/__init__.py b/modules/hidiffusion/__init__.py
index ac7dd9627..2e1a500d9 100644
--- a/modules/hidiffusion/__init__.py
+++ b/modules/hidiffusion/__init__.py
@@ -41,5 +41,5 @@ def apply(p, model_type):
 
 def unapply():
     pipe = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model
-    if hasattr(pipe, 'unet'):
+    if hasattr(pipe, 'unet') and pipe.unet is not None:
         hidiffusion.remove_hidiffusion(pipe)
diff --git a/modules/images.py b/modules/images.py
index bd3cd0c70..d5c2d3130 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -19,6 +19,7 @@ from modules.video import save_video # pylint: disable=unused-import
 
 
 debug = errors.log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug_save = errors.log.trace if os.environ.get('SD_SAVE_DEBUG', None) is not None else lambda *args, **kwargs: None
 try:
     from pi_heif import register_heif_opener
     register_heif_opener()
@@ -26,7 +27,6 @@ except Exception:
     pass
 
 
-
 def sanitize_filename_part(text, replace_spaces=True):
     if text is None:
         return None
@@ -47,8 +47,9 @@ def atomically_save_image():
     while True:
         image, filename, extension, params, exifinfo, filename_txt = save_queue.get()
         shared.state.image_history += 1
-        with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
-            file.write(exifinfo)
+        if len(exifinfo) > 2:
+            with open(paths.params_path, "w", encoding="utf8") as file:
+                file.write(exifinfo)
         fn = filename + extension
         filename = filename.strip()
         if extension[0] != '.': # add dot if missing
@@ -73,6 +74,7 @@ def atomically_save_image():
             pnginfo_data = PngImagePlugin.PngInfo()
             for k, v in params.pnginfo.items():
                 pnginfo_data.add_text(k, str(v))
+            debug_save(f'Save pnginfo: {params.pnginfo.items()}')
             save_args = { 'compress_level': 6, 'pnginfo': pnginfo_data if shared.opts.image_metadata else None }
         elif image_format == 'JPEG':
             if image.mode == 'RGBA':
@@ -82,12 +84,14 @@ def atomically_save_image():
                 image = image.point(lambda p: p * 0.0038910505836576).convert("L")
             save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality }
             if shared.opts.image_metadata:
+                debug_save(f'Save exif: {exifinfo}')
                 save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
         elif image_format == 'WEBP':
             if image.mode == 'I;16':
                 image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
             save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless }
             if shared.opts.image_metadata:
+                debug_save(f'Save exif: {exifinfo}')
                 save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
         elif image_format == 'JXL':
             if image.mode == 'I;16':
@@ -96,10 +100,12 @@ def atomically_save_image():
                 image = image.convert("RGBA")
             save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless }
             if shared.opts.image_metadata:
+                debug_save(f'Save exif: {exifinfo}')
                 save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
         else:
             save_args = { 'quality': shared.opts.jpeg_quality }
         try:
+            debug_save(f'Save args: {save_args}')
             image.save(fn, format=image_format, **save_args)
         except Exception as e:
             shared.log.error(f'Save failed: file="{fn}" format={image_format} args={save_args} {e}')
diff --git a/modules/img2img.py b/modules/img2img.py
index ca71ff0e7..888bf3626 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -32,21 +32,23 @@ def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args)
         inpaint_masks = [f for f in inpaint_masks if filetype.is_image(f)]
         is_inpaint_batch = len(inpaint_masks) > 0
         shared.log.info(f'Process batch: mask folder="{input_dir}" images={len(inpaint_masks)}')
-    save_normally = output_dir == ''
     p.do_not_save_grid = True
-    p.do_not_save_samples = not save_normally
+    p.do_not_save_samples = True
     p.default_prompt = p.prompt
+    if p.n_iter > 1:
+        p.n_iter = 1
+        shared.log.warning(f'Process batch: batch_count={p.n_iter} forced to 1')
     shared.state.job_count = len(image_files) * p.n_iter
     if shared.opts.batch_frame_mode: # SBM Frame mode is on, process each image in batch with same seed
         window_size = p.batch_size
         btcrept = 1
         p.seed = [p.seed] * window_size # SBM MONKEYPATCH: Need to change processing to support a fixed seed value.
         p.subseed = [p.subseed] * window_size # SBM MONKEYPATCH
-        shared.log.info(f"Process batch: inputs={len(image_files)} parallel={window_size} outputs={p.n_iter} per input ")
+        shared.log.info(f"Process batch: inputs={len(image_files)} outputs={p.n_iter}x{len(image_files)} parallel={window_size}")
     else: # SBM Frame mode is off, standard operation of repeating same images with sequential seed.
         window_size = 1
         btcrept = p.batch_size
-        shared.log.info(f"Process batch: inputs={len(image_files)} outputs={p.n_iter * p.batch_size} per input")
+        shared.log.info(f"Process batch: inputs={len(image_files)} outputs={p.n_iter*p.batch_size}x{len(image_files)}")
     for i in range(0, len(image_files), window_size):
         if shared.state.skipped:
             shared.state.skipped = False
@@ -117,8 +119,7 @@ def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args)
                 basename = ''
             if output_dir == '':
                 output_dir = shared.opts.outdir_img2img_samples
-            if not save_normally:
-                os.makedirs(output_dir, exist_ok=True)
+            os.makedirs(output_dir, exist_ok=True)
             geninfo, items = images.read_info_from_image(image)
             for k, v in items.items():
                 image.info[k] = v
diff --git a/modules/infiniteyou/pipeline_infu_flux.py b/modules/infiniteyou/pipeline_infu_flux.py
index 2b4762d97..fc8b84647 100644
--- a/modules/infiniteyou/pipeline_infu_flux.py
+++ b/modules/infiniteyou/pipeline_infu_flux.py
@@ -151,7 +151,7 @@ class InfUFluxPipeline:
         local_path = snapshot_download(repo_id='ByteDance/InfiniteYou', cache_dir=shared.opts.hfcache_dir)
         infiniteyou_path = os.path.join(local_path, f'infu_flux_{infu_flux_version}', model_version)
         infusenet_path = os.path.join(infiniteyou_path, 'InfuseNetModel')
-        quant_args = model_quant.create_config(module='ControlNet')
+        quant_args = model_quant.create_config(module='Control')
         shared.log.debug(f'InfiniteYou: fn="{infusenet_path}" load infusenet')
         self.infusenet = FluxControlNetModel.from_pretrained(
             infusenet_path,
diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py
index d3487fefd..1f0391295 100644
--- a/modules/intel/ipex/diffusers.py
+++ b/modules/intel/ipex/diffusers.py
@@ -26,7 +26,7 @@ class FluxPosEmbed(torch.nn.Module):
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
-        pos = ids.float()
+        pos = ids.to(dtype=torch.float32)
         for i in range(n_axes):
             cos, sin = diffusers.models.embeddings.get_1d_rotary_pos_embed(
                 self.axes_dim[i],
@@ -99,12 +99,12 @@ def apply_rotary_emb(x, freqs_cis, use_real: bool = True, use_real_unbind_dim: i
         else:
             raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
 
-        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+        out = (x.to(dtype=torch.float32) * cos + x_rotated.to(dtype=torch.float32) * sin).to(x.dtype)
         return out
     else:
         # used for lumina
         # force cpu with Alchemist
-        x_rotated = torch.view_as_complex(x.to("cpu").float().reshape(*x.shape[:-1], -1, 2))
+        x_rotated = torch.view_as_complex(x.to("cpu").to(dtype=torch.float32).reshape(*x.shape[:-1], -1, 2))
         freqs_cis = freqs_cis.to("cpu").unsqueeze(2)
         x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
         return x_out.type_as(x).to(x.device)
@@ -122,5 +122,6 @@ def ipex_diffusers(device_supports_fp64=False):
         diffusers.models.embeddings.apply_rotary_emb = apply_rotary_emb
         diffusers.models.transformers.transformer_flux.FluxPosEmbed = FluxPosEmbed
         diffusers.models.transformers.transformer_lumina2.apply_rotary_emb = apply_rotary_emb
-        diffusers.models.controlnets.controlnet_flux.FluxPosEmbed = FluxPosEmbed
         diffusers.models.transformers.transformer_hidream_image.rope = hidream_rope
+        diffusers.models.transformers.transformer_chroma.FluxPosEmbed = FluxPosEmbed
+        diffusers.models.controlnets.controlnet_flux.FluxPosEmbed = FluxPosEmbed
diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py
index d81d7b05c..e2a04a662 100644
--- a/modules/intel/ipex/hijacks.py
+++ b/modules/intel/ipex/hijacks.py
@@ -254,8 +254,15 @@ torch.Tensor.original_Tensor_to = torch.Tensor.to
 @wraps(torch.Tensor.to)
 def Tensor_to(self, device=None, *args, **kwargs):
     if check_cuda(device):
+        if not device_supports_fp64 and kwargs.get("dtype", None) == torch.float64:
+            kwargs["dtype"] = torch.float32
         return self.original_Tensor_to(return_xpu(device), *args, **kwargs)
     else:
+        if not device_supports_fp64:
+            if kwargs.get("dtype", None) == torch.float64 and ((device is None and self.device.type == "xpu") or (device is not None and torch.device(device).type == "xpu")):
+                kwargs["dtype"] = torch.float32
+            elif device == torch.float64 and self.device.type == "xpu":
+                device = torch.float32
         return self.original_Tensor_to(device, *args, **kwargs)
 
 original_Tensor_cuda = torch.Tensor.cuda
diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py
index bb5ee9c64..26ac004e4 100644
--- a/modules/intel/openvino/__init__.py
+++ b/modules/intel/openvino/__init__.py
@@ -24,11 +24,13 @@ from modules import shared, devices, sd_models
 
 # importing openvino.runtime forces DeprecationWarning to "always"
 # And Intel's own libs (NNCF) imports the deprecated module
-# Reset the warnings back to ignore:
+# Don't allow openvino to override warning filters:
 try:
     import warnings
+    filterwarnings = warnings.filterwarnings
+    warnings.filterwarnings = lambda *args, **kwargs: None
     import openvino.runtime # pylint: disable=unused-import
-    warnings.filterwarnings(action="ignore", category=DeprecationWarning)
+    warnings.filterwarnings = filterwarnings
 except Exception:
     pass
 
diff --git a/modules/interposer.py b/modules/interposer.py
index 2bf272706..c248c4f00 100644
--- a/modules/interposer.py
+++ b/modules/interposer.py
@@ -101,7 +101,7 @@ def map_model_name(name: str):
         return 'xl'
     if name == 'sd3':
         return 'v3'
-    if name == 'f1':
+    if name in ['f1', 'chroma']:
         return 'fx'
     return name
 
diff --git a/modules/interrogate/joycaption.py b/modules/interrogate/joycaption.py
index cc316341a..dc5e07213 100644
--- a/modules/interrogate/joycaption.py
+++ b/modules/interrogate/joycaption.py
@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 import torch
 from transformers import AutoProcessor, LlavaForConditionalGeneration
-from modules import shared, devices
+from modules import shared, devices, sd_models, model_quant
 
 
 """
@@ -58,13 +58,24 @@ opts = JoyOptions()
 
 
 @torch.no_grad()
-def predict(question: str, image):
+def predict(question: str, image, vqa_model: str = None) -> str:
     global llava_model, processor # pylint: disable=global-statement
     opts.max_new_tokens = shared.opts.interrogate_vlm_max_length
+    if vqa_model is not None and opts.repo != vqa_model:
+        opts.repo = vqa_model
+        llava_model = None
     if llava_model is None:
         shared.log.info(f'Interrogate: type=vlm model="JoyCaption" {str(opts)}')
+
         processor = AutoProcessor.from_pretrained(opts.repo)
-        llava_model = LlavaForConditionalGeneration.from_pretrained(opts.repo, torch_dtype=devices.dtype, device_map="auto", cache_dir=shared.opts.hfcache_dir)
+        quant_args = model_quant.create_config(module='LLM')
+        llava_model = LlavaForConditionalGeneration.from_pretrained(
+            opts.repo,
+            torch_dtype=devices.dtype,
+            device_map="auto",
+            cache_dir=shared.opts.hfcache_dir,
+            **quant_args,
+        )
         llava_model.eval()
 
     if len(question) < 2:
@@ -77,7 +88,7 @@ def predict(question: str, image):
     convo_string = processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
     inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(devices.device) # Process the inputs
     inputs['pixel_values'] = inputs['pixel_values'].to(devices.dtype)
-    llava_model = llava_model.to(devices.device)
+    sd_models.move_model(llava_model, devices.device)
     with devices.inference_context():
         generate_ids = llava_model.generate( # Generate the captions
             **inputs,
@@ -94,6 +105,6 @@ def predict(question: str, image):
         )[0]
         generate_ids = generate_ids[inputs['input_ids'].shape[1]:] # Trim off the prompt
         caption = processor.tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) # Decode the caption
-    llava_model = llava_model.to(devices.cpu)
+    sd_models.move_model(llava_model, devices.cpu, force=True)
     caption = caption.replace('\n\n', '\n').strip()
     return caption
diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py
index 32820a243..d2aeab61a 100644
--- a/modules/interrogate/vqa.py
+++ b/modules/interrogate/vqa.py
@@ -7,12 +7,13 @@ import torch
 import transformers
 import transformers.dynamic_module_utils
 from PIL import Image
-from modules import shared, devices, errors, sd_models
+from modules import shared, devices, errors, sd_models, model_quant
 
 
 processor = None
 model = None
 loaded: str = None
+quant_args = {}
 vlm_models = {
     "Microsoft Florence 2 Base": "microsoft/Florence-2-base", # 0.5GB
     "Microsoft Florence 2 Large": "microsoft/Florence-2-large", # 1.5GB
@@ -38,7 +39,8 @@ vlm_models = {
     "ToriiGate 0.4 2B": "Minthy/ToriiGate-v0.4-2B",
     "ToriiGate 0.4 7B": "Minthy/ToriiGate-v0.4-7B",
     "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB
-    "JoyCaption": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB
+    "JoyCaption Alpha": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB
+    "JoyCaption Beta": "fancyfeast/llama-joycaption-beta-one-hf-llava", # 17.4GB
     "JoyTag": "fancyfeast/joytag", # 0.7GB
     "AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B",
     "AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B",
@@ -120,9 +122,10 @@ def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str
         model = None
         model = transformers.Qwen2VLForConditionalGeneration.from_pretrained(
             repo,
-            cache_dir=shared.opts.hfcache_dir
+            torch_dtype=devices.dtype,
+            cache_dir=shared.opts.hfcache_dir,
+            **quant_args,
         )
-        model = model.to(devices.device, devices.dtype)
         processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
         devices.torch_gc()
@@ -165,8 +168,12 @@ def gemma(question: str, image: Image.Image, repo: str = None, system_prompt: st
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
         model = None
-        model = transformers.Gemma3ForConditionalGeneration.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
-        model = model.to(devices.device, devices.dtype)
+        model = transformers.Gemma3ForConditionalGeneration.from_pretrained(
+            repo,
+            torch_dtype=devices.dtype,
+            cache_dir=shared.opts.hfcache_dir,
+            **quant_args,
+        )
         processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
         devices.torch_gc()
@@ -216,7 +223,6 @@ def paligemma(question: str, image: Image.Image, repo: str = None):
             cache_dir=shared.opts.hfcache_dir,
             torch_dtype=devices.dtype,
         )
-        model = model.to(devices.device, devices.dtype)
         loaded = repo
         devices.torch_gc()
     sd_models.move_model(model, devices.device)
@@ -250,7 +256,6 @@ def ovis(question: str, image: Image.Image, repo: str = None):
             trust_remote_code=True,
             cache_dir=shared.opts.hfcache_dir,
         )
-        model = model.to(devices.device, devices.dtype)
         loaded = repo
         devices.torch_gc()
     sd_models.move_model(model, devices.device)
@@ -291,8 +296,8 @@ def smol(question: str, image: Image.Image, repo: str = None, system_prompt: str
             cache_dir=shared.opts.hfcache_dir,
             torch_dtype=devices.dtype,
             _attn_implementation="eager",
+            **quant_args,
             )
-        model.to(devices.device, devices.dtype)
         processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
         devices.torch_gc()
@@ -330,9 +335,9 @@ def git(question: str, image: Image.Image, repo: str = None):
         model = None
         model = transformers.GitForCausalLM.from_pretrained(
             repo,
+            torch_dtype=devices.dtype,
             cache_dir=shared.opts.hfcache_dir,
         )
-        model.to(devices.device, devices.dtype)
         processor = transformers.GitProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
         devices.torch_gc()
@@ -358,9 +363,9 @@ def blip(question: str, image: Image.Image, repo: str = None):
         model = None
         model = transformers.BlipForQuestionAnswering.from_pretrained(
             repo,
+            torch_dtype=devices.dtype,
             cache_dir=shared.opts.hfcache_dir,
         )
-        model.to(devices.device, devices.dtype)
         processor = transformers.BlipProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
         devices.torch_gc()
@@ -380,9 +385,9 @@ def vilt(question: str, image: Image.Image, repo: str = None):
         model = None
         model = transformers.ViltForQuestionAnswering.from_pretrained(
             repo,
+            torch_dtype=devices.dtype,
             cache_dir=shared.opts.hfcache_dir,
         )
-        model.to(devices.device)
         processor = transformers.ViltProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
         devices.torch_gc()
@@ -404,9 +409,9 @@ def pix(question: str, image: Image.Image, repo: str = None):
         model = None
         model = transformers.Pix2StructForConditionalGeneration.from_pretrained(
             repo,
+            torch_dtype=devices.dtype,
             cache_dir=shared.opts.hfcache_dir,
         )
-        model.to(devices.device)
         processor = transformers.Pix2StructProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
         devices.torch_gc()
@@ -428,20 +433,30 @@ def moondream(question: str, image: Image.Image, repo: str = None):
         model = None
         model = transformers.AutoModelForCausalLM.from_pretrained(
             repo,
-            revision="2024-08-26",
+            revision="2025-06-21",
             trust_remote_code=True,
-            cache_dir=shared.opts.hfcache_dir
+            torch_dtype=devices.dtype,
+            cache_dir=shared.opts.hfcache_dir,
         )
         processor = transformers.AutoTokenizer.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
-        model.to(devices.device, devices.dtype)
         model.eval()
         devices.torch_gc()
     sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
     encoded = model.encode_image(image)
     with devices.inference_context():
-        response = model.answer_question(encoded, question, processor)
+        if question == 'CAPTION':
+            response = model.caption(image, length="short")['caption']
+        elif question == 'DETAILED CAPTION':
+            response = model.caption(image, length="normal")['caption']
+        elif question == 'MORE DETAILED CAPTION':
+            response = model.caption(image, length="long")['caption']
+        else:
+            response = model.answer_question(encoded, question, processor)['answer']
+        # model.detect(image, "face")
+        # model.point(image, "person")
+        # model.detect_gaze(image)
     return response
 
 
@@ -464,12 +479,13 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str
             repo,
             trust_remote_code=True,
             revision=revision,
+            torch_dtype=devices.dtype,
             cache_dir=shared.opts.hfcache_dir,
+            **quant_args,
         )
         processor = transformers.AutoProcessor.from_pretrained(repo, trust_remote_code=True, revision=revision, cache_dir=shared.opts.hfcache_dir)
         transformers.dynamic_module_utils.get_imports = _get_imports
         loaded = repo
-        model.to(devices.device, devices.dtype)
         model.eval()
         devices.torch_gc()
     sd_models.move_model(model, devices.device)
@@ -501,7 +517,6 @@ def sa2(question: str, image: Image.Image, repo: str = None):
             low_cpu_mem_usage=True,
             use_flash_attn=False,
             trust_remote_code=True)
-        model = model.to(devices.device, devices.dtype)
         model = model.eval()
         processor = transformers.AutoTokenizer.from_pretrained(
             repo,
@@ -528,9 +543,11 @@ def sa2(question: str, image: Image.Image, repo: str = None):
 
 
 def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:Image.Image=None, model_name:str=None, quiet:bool=False):
+    global quant_args # pylint: disable=global-statement
     if not quiet:
         shared.state.begin('Interrogate')
     t0 = time.time()
+    quant_args = model_quant.create_config(module='LLM')
     model_name = model_name or shared.opts.interrogate_vlm_model
     if isinstance(image, list):
         image = image[0] if len(image) > 0 else None
@@ -549,8 +566,10 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:
     if shared.native and shared.sd_loaded:
         from modules.sd_models import apply_balanced_offload # prevent circular import
         apply_balanced_offload(shared.sd_model)
+
     from modules import modelloader
     modelloader.hf_login()
+
     try:
         if model_name is None:
             shared.log.error(f'Interrogate: type=vlm model="{model_name}" no model selected')
@@ -562,6 +581,7 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:
         if image is None:
             shared.log.error(f'Interrogate: type=vlm model="{model_name}" no input image')
             return ''
+
         if 'git' in vqa_model.lower():
             answer = git(question, image, vqa_model)
         elif 'vilt' in vqa_model.lower():
@@ -583,7 +603,7 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:
             answer = joytag.predict(image)
         elif 'joycaption' in vqa_model.lower():
             from modules.interrogate import joycaption
-            answer = joycaption.predict(question, image)
+            answer = joycaption.predict(question, image, vqa_model)
         elif 'deepseek' in vqa_model.lower():
             from modules.interrogate import deepseek
             answer = deepseek.predict(question, image, vqa_model)
@@ -600,9 +620,10 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:
     except Exception as e:
         errors.display(e, 'VQA')
         answer = 'error'
+
     if shared.opts.interrogate_offload and model is not None:
-        sd_models.move_model(model, devices.cpu)
-    devices.torch_gc()
+        sd_models.move_model(model, devices.cpu, force=True)
+    devices.torch_gc(force=True)
     answer = clean(answer, question)
     t1 = time.time()
     if not quiet:
diff --git a/modules/ipadapter.py b/modules/ipadapter.py
index a03381a1b..3b2ce3ea9 100644
--- a/modules/ipadapter.py
+++ b/modules/ipadapter.py
@@ -146,7 +146,7 @@ def unapply(pipe, unload: bool = False): # pylint: disable=arguments-differ
             if unload:
                 shared.log.debug('IP adapter unload')
                 pipe.unload_ip_adapter()
-        if hasattr(pipe, 'unet'):
+        if hasattr(pipe, 'unet') and pipe.unet is not None:
             module = pipe.unet
         elif hasattr(pipe, 'transformer'):
             module = pipe.transformer
diff --git a/modules/loader.py b/modules/loader.py
index 2574c90b4..674167d20 100644
--- a/modules/loader.py
+++ b/modules/loader.py
@@ -35,9 +35,14 @@ warnings.filterwarnings(action="ignore", category=DeprecationWarning)
 warnings.filterwarnings(action="ignore", category=FutureWarning)
 warnings.filterwarnings(action="ignore", category=UserWarning, module="torchvision")
 try:
+    import torch._logging # pylint: disable=ungrouped-imports
+    torch._logging._internal.DEFAULT_LOG_LEVEL = logging.ERROR # pylint: disable=protected-access
     torch._logging.set_logs(all=logging.ERROR, bytecode=False, aot_graphs=False, aot_joint_graph=False, ddp_graphs=False, graph=False, graph_code=False, graph_breaks=False, graph_sizes=False, guards=False, recompiles=False, recompiles_verbose=False, trace_source=False, trace_call=False, trace_bytecode=False, output_code=False, kernel_code=False, schedule=False, perf_hints=False, post_grad_graphs=False, onnx_diagnostics=False, fusion=False, overlap=False, export=None, modules=None, cudagraphs=False, sym_node=False, compiled_autograd_verbose=False) # pylint: disable=protected-access
-except Exception:
-    pass
+    import torch._dynamo
+    torch._dynamo.config.verbose = False # pylint: disable=protected-access
+    torch._dynamo.config.suppress_errors = True # pylint: disable=protected-access
+except Exception as e:
+    errors.log.warning(f'Torch logging: {e}')
 if ".dev" in torch.__version__ or "+git" in torch.__version__:
     torch.__long_version__ = torch.__version__
     torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0)
diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py
index 40982caac..6f5e1b42a 100644
--- a/modules/lora/extra_networks_lora.py
+++ b/modules/lora/extra_networks_lora.py
@@ -141,10 +141,14 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
     def changed(self, requested: List[str], include: List[str], exclude: List[str]):
         if shared.opts.lora_force_reload:
             return True
-        sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+        sd_model = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model
         if not hasattr(sd_model, 'loaded_loras'):
             sd_model.loaded_loras = {}
-        key = f'{",".join(include)}:{",".join(exclude)}'
+        if include is None or len(include) == 0:
+            include = ['all']
+        if exclude is None or len(exclude) == 0:
+            exclude = ['none']
+        key = f'include={",".join(include)}:exclude={",".join(exclude)}'
         loaded = sd_model.loaded_loras.get(key, [])
         debug_log(f'Network load: type=LoRA key="{key}" requested={requested} loaded={loaded}')
         if len(requested) != len(loaded):
diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py
index b971541cd..2fcea174c 100644
--- a/modules/lora/lora_apply.py
+++ b/modules/lora/lora_apply.py
@@ -79,7 +79,9 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
             continue
         try:
             t0 = time.time()
-            if hasattr(self, "sdnq_dequantizer"):
+            if hasattr(self, "sdnq_dequantizer_backup"):
+                weight = self.sdnq_dequantizer_backup.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul)
+            elif hasattr(self, "sdnq_dequantizer"):
                 weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul)
             else:
                 weight = self.weight.to(devices.device) # must perform calc on gpu due to performance
@@ -228,6 +230,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
             self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False)
             if hasattr(self, "sdnq_dequantizer_backup"):
                 self.sdnq_dequantizer = self.sdnq_dequantizer_backup.to(device)
+                del self.sdnq_dequantizer_backup
 
     if bias_backup is not None:
         self.bias = None
diff --git a/modules/lora/lora_extract.py b/modules/lora/lora_extract.py
index 4bd3a3adb..27bfb5c20 100644
--- a/modules/lora/lora_extract.py
+++ b/modules/lora/lora_extract.py
@@ -107,6 +107,9 @@ def make_meta(fn, maxrank, rank_ratio):
     elif shared.sd_model_type == "f1":
         meta["model_spec.architecture"] = "flux-1-dev/lora"
         meta["ss_base_model_version"] = "flux1"
+    elif shared.sd_model_type == "chroma":
+        meta["model_spec.architecture"] = "chroma/lora"
+        meta["ss_base_model_version"] = "chroma"
     elif shared.sd_model_type == "sc":
         meta["model_spec.architecture"] = "stable-cascade-v1-prior/lora"
     return meta
diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py
index f02e4a274..cec3e2542 100644
--- a/modules/lora/lora_load.py
+++ b/modules/lora/lora_load.py
@@ -85,7 +85,7 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
     net = network.Network(name, network_on_disk)
     net.mtime = os.path.getmtime(network_on_disk.filename)
     state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
-    if shared.sd_model_type == 'f1': # if kohya flux lora, convert state_dict
+    if shared.sd_model_type in ['f1', 'chroma']: # if kohya flux lora, convert state_dict
         state_dict = lora_convert._convert_kohya_flux_lora_to_diffusers(state_dict) or state_dict # pylint: disable=protected-access
     if shared.sd_model_type == 'sd3': # if kohya flux lora, convert state_dict
         try:
@@ -148,8 +148,8 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
             if net_module is not None:
                 network_types.append(nettype.__class__.__name__)
                 break
-            module_errors += 1
         if net_module is None:
+            module_errors += 1
             if l.debug:
                 shared.log.error(f'LoRA unhandled: name={name} key={key} weights={weights.w.keys()}')
         else:
diff --git a/modules/lora/network.py b/modules/lora/network.py
index f6d93009c..94c9cecf4 100644
--- a/modules/lora/network.py
+++ b/modules/lora/network.py
@@ -18,6 +18,7 @@ class SdVersion(enum.Enum):
     SC = 5
     F1 = 6
     HV = 7
+    CHROMA = 8
 
 
 class NetworkOnDisk:
@@ -59,6 +60,8 @@ class NetworkOnDisk:
             return 'f1'
         if base.startswith("hunyuan_video"):
             return 'hv'
+        if base.startswith("chroma"):
+            return 'chroma'
 
         if arch.startswith("stable-diffusion-v1"):
             return 'sd1'
@@ -70,6 +73,8 @@ class NetworkOnDisk:
             return 'f1'
         if arch.startswith("hunyuan-video"):
             return 'hv'
+        if arch.startswith("chroma"):
+            return 'chroma'
 
         if "v1-5" in str(self.metadata.get('ss_sd_model_name', "")):
             return 'sd1'
@@ -79,6 +84,8 @@ class NetworkOnDisk:
             return 'f1'
         if 'xl' in self.name.lower():
             return 'xl'
+        if 'chroma' in self.name.lower():
+            return 'chroma'
 
         return ''
 
@@ -90,6 +97,27 @@ class NetworkOnDisk:
         if not self.hash:
             self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '')
 
+    def get_info(self):
+        data = {}
+        if shared.cmd_opts.no_metadata:
+            return data
+        if self.filename is not None:
+            fn = os.path.splitext(self.filename)[0] + '.json'
+            if os.path.exists(fn):
+                data = shared.readfile(fn, silent=True)
+                if type(data) is list:
+                    data = data[0]
+        return data
+
+    def get_desc(self):
+        if shared.cmd_opts.no_metadata:
+            return None
+        if self.filename is not None:
+            fn = os.path.splitext(self.filename)[0] + '.txt'
+            if os.path.exists(fn):
+                return shared.readfile(fn, silent=True)
+        return None
+
     def get_alias(self):
         if shared.opts.lora_preferred_name == "filename":
             return self.name
@@ -124,7 +152,10 @@ class NetworkModule:
         self.sd_key = weights.sd_key
         self.sd_module = weights.sd_module
         if hasattr(self.sd_module, 'weight'):
-            self.shape = self.sd_module.weight.shape
+            if hasattr(self.sd_module, "sdnq_dequantizer"):
+                self.shape = self.sd_module.sdnq_dequantizer.original_shape
+            else:
+                self.shape = self.sd_module.weight.shape
         self.dim = None
         self.bias = weights.w.get("bias")
         self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None
diff --git a/modules/mit_nunchaku.py b/modules/mit_nunchaku.py
index 9c0500790..798de661e 100644
--- a/modules/mit_nunchaku.py
+++ b/modules/mit_nunchaku.py
@@ -1,13 +1,10 @@
 # MIT-Han-Lab Nunchaku: <https://github.com/mit-han-lab/nunchaku>
-# TODO nunchaku: cache-dir for transformer and t5 loader
-# TODO nunchaku: batch support
-
 
 from installer import log, pip
 from modules import devices
 
 
-ver = '0.2.0'
+ver = '0.3.1'
 ok = False
 
 
diff --git a/modules/model_chroma.py b/modules/model_chroma.py
new file mode 100644
index 000000000..9333ccc30
--- /dev/null
+++ b/modules/model_chroma.py
@@ -0,0 +1,299 @@
+import os
+import json
+import torch
+import diffusers
+import transformers
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download, auth_check
+from modules import shared, errors, devices, modelloader, sd_models, sd_unet, model_te, model_quant, sd_hijack_te
+
+
+debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
+
+
+def load_chroma_quanto(checkpoint_info):
+    transformer, text_encoder = None, None
+    quanto = model_quant.load_quanto('Load model: type=Chroma')
+
+    if isinstance(checkpoint_info, str):
+        repo_path = checkpoint_info
+    else:
+        repo_path = checkpoint_info.path
+
+    try:
+        quantization_map = os.path.join(repo_path, "transformer", "quantization_map.json")
+        debug(f'Load model: type=Chroma quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="transformer"')
+        if not os.path.exists(quantization_map):
+            repo_id = sd_models.path_to_repo(checkpoint_info.name)
+            quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir)
+        with open(quantization_map, "r", encoding='utf8') as f:
+            quantization_map = json.load(f)
+        state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors"))
+        dtype = state_dict['context_embedder.bias'].dtype
+        with torch.device("meta"):
+            transformer = diffusers.ChromaTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype)
+        quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu"))
+        if shared.opts.diffusers_eval:
+            transformer.eval()
+        transformer_dtype = transformer.dtype
+        if transformer_dtype != devices.dtype:
+            try:
+                transformer = transformer.to(dtype=devices.dtype)
+            except Exception:
+                shared.log.error(f"Load model: type=Chroma Failed to cast transformer to {devices.dtype}, set dtype to {transformer_dtype}")
+    except Exception as e:
+        shared.log.error(f"Load model: type=Chroma failed to load Quanto transformer: {e}")
+        if debug:
+            errors.display(e, 'Chroma Quanto:')
+
+    try:
+        quantization_map = os.path.join(repo_path, "text_encoder", "quantization_map.json")
+        debug(f'Load model: type=Chroma quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="text_encoder"')
+        if not os.path.exists(quantization_map):
+            repo_id = sd_models.path_to_repo(checkpoint_info.name)
+            quantization_map = hf_hub_download(repo_id, subfolder='text_encoder', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir)
+        with open(quantization_map, "r", encoding='utf8') as f:
+            quantization_map = json.load(f)
+        with open(os.path.join(repo_path, "text_encoder", "config.json"), encoding='utf8') as f:
+            t5_config = transformers.T5Config(**json.load(f))
+        state_dict = load_file(os.path.join(repo_path, "text_encoder", "model.safetensors"))
+        dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype
+        with torch.device("meta"):
+            text_encoder = transformers.T5EncoderModel(t5_config).to(dtype=dtype)
+        quanto.requantize(text_encoder, state_dict, quantization_map, device=torch.device("cpu"))
+        if shared.opts.diffusers_eval:
+            text_encoder.eval()
+        text_encoder_dtype = text_encoder.dtype
+        if text_encoder_dtype != devices.dtype:
+            try:
+                text_encoder = text_encoder.to(dtype=devices.dtype)
+            except Exception:
+                shared.log.error(f"Load model: type=Chroma Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_dtype}")
+    except Exception as e:
+        shared.log.error(f"Load model: type=Chroma failed to load Quanto text encoder: {e}")
+        if debug:
+            errors.display(e, 'Chroma Quanto:')
+
+    return transformer, text_encoder
+
+
+def load_chroma_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unused-argument
+    transformer, text_encoder = None, None
+    if isinstance(checkpoint_info, str):
+        repo_path = checkpoint_info
+    else:
+        repo_path = checkpoint_info.path
+    model_quant.load_bnb('Load model: type=Chroma')
+    quant = model_quant.get_quant(repo_path)
+    try:
+        # we ignore the distilled guidance layer because it degrades quality too much
+        # see: https://github.com/huggingface/diffusers/pull/11698#issuecomment-2969717180 for more details
+        if quant == 'fp8':
+            quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["distilled_guidance_layer"], bnb_4bit_compute_dtype=devices.dtype)
+            debug(f'Quantization: {quantization_config}')
+            transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
+        elif quant == 'fp4':
+            quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, llm_int8_skip_modules=["distilled_guidance_layer"], bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'fp4')
+            debug(f'Quantization: {quantization_config}')
+            transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
+        elif quant == 'nf4':
+            quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, llm_int8_skip_modules=["distilled_guidance_layer"], bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'nf4')
+            debug(f'Quantization: {quantization_config}')
+            transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config)
+        else:
+            transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config)
+    except Exception as e:
+        shared.log.error(f"Load model: type=Chroma failed to load BnB transformer: {e}")
+        transformer, text_encoder = None, None
+        if debug:
+            errors.display(e, 'Chroma:')
+    return transformer, text_encoder
+
+
+def load_quants(kwargs, pretrained_model_name_or_path, cache_dir, allow_quant):
+    try:
+        if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'):
+            raise NotImplementedError('Nunchaku does not support Chroma Model yet. See https://github.com/mit-han-lab/nunchaku/issues/167')
+        elif 'transformer' not in kwargs and model_quant.check_quant('Model'):
+            quant_args = model_quant.create_config(allow=allow_quant, module='Model', modules_to_not_convert=["distilled_guidance_layer"])
+            if quant_args:
+                if os.path.isfile(pretrained_model_name_or_path):
+                    kwargs['transformer'] = diffusers.ChromaTransformer2DModel.from_single_file(pretrained_model_name_or_path, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
+                else:
+                    kwargs['transformer'] = diffusers.ChromaTransformer2DModel.from_pretrained(pretrained_model_name_or_path, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
+        if 'text_encoder' not in kwargs and model_quant.check_nunchaku('TE'):
+            import nunchaku
+            nunchaku_precision = nunchaku.utils.get_precision()
+            nunchaku_repo = 'mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors'
+            shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}')
+            kwargs['text_encoder'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
+        elif 'text_encoder' not in kwargs and model_quant.check_quant('TE'):
+            quant_args = model_quant.create_config(allow=allow_quant, module='TE')
+            if quant_args:
+                if os.path.isfile(pretrained_model_name_or_path):
+                    kwargs['text_encoder'] = transformers.T5EncoderModel.from_single_file(pretrained_model_name_or_path, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
+                else:
+                    kwargs['text_encoder'] = transformers.T5EncoderModel.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
+    except Exception as e:
+        shared.log.error(f'Quantization: {e}')
+        errors.display(e, 'Quantization:')
+    return kwargs
+
+
+def load_transformer(file_path): # triggered by opts.sd_unet change
+    if file_path is None or not os.path.exists(file_path):
+        return None
+    transformer = None
+    quant = model_quant.get_quant(file_path)
+    diffusers_load_config = {
+        "low_cpu_mem_usage": True,
+        "torch_dtype": devices.dtype,
+        "cache_dir": shared.opts.hfcache_dir,
+    }
+    if quant is not None and quant != 'none':
+        shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} prequant={quant} dtype={devices.dtype}')
+    if 'gguf' in file_path.lower():
+        from modules import ggml
+        _transformer = ggml.load_gguf(file_path, cls=diffusers.ChromaTransformer2DModel, compute_dtype=devices.dtype)
+        if _transformer is not None:
+            transformer = _transformer
+    elif quant == 'qint8' or quant == 'qint4':
+        _transformer, _text_encoder = load_chroma_quanto(file_path)
+        if _transformer is not None:
+            transformer = _transformer
+    elif quant == 'fp8' or quant == 'fp4' or quant == 'nf4':
+        _transformer, _text_encoder = load_chroma_bnb(file_path, diffusers_load_config)
+        if _transformer is not None:
+            transformer = _transformer
+    else:
+        quant_args = model_quant.create_config(module='Model', modules_to_not_convert=["distilled_guidance_layer"])
+        if quant_args:
+            shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=torchao dtype={devices.dtype}')
+            transformer = diffusers.ChromaTransformer2DModel.from_single_file(file_path, **diffusers_load_config, **quant_args)
+            if transformer is not None:
+                return transformer
+        shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=none dtype={devices.dtype}')
+        # TODO chroma transformer from-single-file with quant
+        # shared.log.warning('Load module: type=UNet/Transformer does not support load-time quantization')
+        # transformer = diffusers.ChromaTransformer2DModel.from_single_file(file_path, **diffusers_load_config)
+    if transformer is None:
+        shared.log.error('Failed to load UNet model')
+        shared.opts.sd_unet = 'Default'
+    return transformer
+
+
+def load_chroma(checkpoint_info, diffusers_load_config): # triggered by opts.sd_checkpoint change
+    fn = checkpoint_info.path
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    login = modelloader.hf_login()
+    try:
+        auth_check(repo_id)
+    except Exception as e:
+        repo_id = None
+        if not os.path.exists(fn):
+            shared.log.error(f'Load model: repo="{repo_id}" login={login} {e}')
+            return None
+
+    prequantized = model_quant.get_quant(checkpoint_info.path)
+    shared.log.debug(f'Load model: type=Chroma model="{checkpoint_info.name}" repo={repo_id or "none"} unet="{shared.opts.sd_unet}" te="{shared.opts.sd_text_encoder}" vae="{shared.opts.sd_vae}" quant={prequantized} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
+    debug(f'Load model: type=Chroma config={diffusers_load_config}')
+
+    transformer = None
+    text_encoder = None
+    vae = None
+
+    # unload current model
+    sd_models.unload_model_weights()
+    shared.sd_model = None
+    devices.torch_gc(force=True)
+
+    if shared.opts.teacache_enabled:
+        from modules import teacache
+        shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.ChromaTransformer2DModel.__name__}')
+        diffusers.ChromaTransformer2DModel.forward = teacache.teacache_chroma_forward # patch must be done before transformer is loaded
+
+    # load overrides if any
+    if shared.opts.sd_unet != 'Default':
+        try:
+            debug(f'Load model: type=Chroma unet="{shared.opts.sd_unet}"')
+            transformer = load_transformer(sd_unet.unet_dict[shared.opts.sd_unet])
+            if transformer is None:
+                shared.opts.sd_unet = 'Default'
+                sd_unet.failed_unet.append(shared.opts.sd_unet)
+        except Exception as e:
+            shared.log.error(f"Load model: type=Chroma failed to load UNet: {e}")
+            shared.opts.sd_unet = 'Default'
+            if debug:
+                errors.display(e, 'Chroma UNet:')
+    if shared.opts.sd_text_encoder != 'Default':
+        try:
+            debug(f'Load model: type=Chroma te="{shared.opts.sd_text_encoder}"')
+            from modules.model_te import load_t5
+            text_encoder = load_t5(name=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
+        except Exception as e:
+            shared.log.error(f"Load model: type=Chroma failed to load T5: {e}")
+            shared.opts.sd_text_encoder = 'Default'
+            if debug:
+                errors.display(e, 'Chroma T5:')
+    if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
+        try:
+            debug(f'Load model: type=Chroma vae="{shared.opts.sd_vae}"')
+            from modules import sd_vae
+            # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override')
+            vae_file = sd_vae.vae_dict[shared.opts.sd_vae]
+            if os.path.exists(vae_file):
+                vae_config = os.path.join('configs', 'chroma', 'vae', 'config.json')
+                vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config)
+        except Exception as e:
+            shared.log.error(f"Load model: type=Chroma failed to load VAE: {e}")
+            shared.opts.sd_vae = 'Default'
+            if debug:
+                errors.display(e, 'Chroma VAE:')
+
+    # initialize pipeline with pre-loaded components
+    kwargs = {}
+    if transformer is not None:
+        kwargs['transformer'] = transformer
+        sd_unet.loaded_unet = shared.opts.sd_unet
+    if text_encoder is not None:
+        kwargs['text_encoder'] = text_encoder
+        model_te.loaded_te = shared.opts.sd_text_encoder
+    if vae is not None:
+        kwargs['vae'] = vae
+
+    # TODO add ChromaFillPipeline, ChromaControlPipeline, ChromaImg2ImgPipeline etc when available
+    # TODO Chroma will support inpainting *after* its training has finished: https://huggingface.co/lodestones/Chroma/discussions/28#6826dd2ed86f53ff983add5c
+    cls = diffusers.ChromaPipeline
+    shared.log.debug(f'Load model: type=Chroma cls={cls.__name__} preloaded={list(kwargs)} revision={diffusers_load_config.get("revision", None)}')
+    for c in kwargs:
+        if getattr(kwargs[c], 'quantization_method', None) is not None or getattr(kwargs[c], 'gguf', None) is not None:
+            shared.log.debug(f'Load model: type=Chroma component={c} dtype={kwargs[c].dtype} quant={getattr(kwargs[c], "quantization_method", None) or getattr(kwargs[c], "gguf", None)}')
+        if kwargs[c].dtype == torch.float32 and devices.dtype != torch.float32:
+            try:
+                kwargs[c] = kwargs[c].to(dtype=devices.dtype)
+                shared.log.warning(f'Load model: type=Chroma component={c} dtype={kwargs[c].dtype} cast dtype={devices.dtype} recast')
+            except Exception:
+                pass
+
+    allow_quant = 'gguf' not in (sd_unet.loaded_unet or '') and (prequantized is None or prequantized == 'none')
+    if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
+        kwargs = load_quants(kwargs, repo_id or fn, cache_dir=shared.opts.diffusers_dir, allow_quant=allow_quant)
+    # kwargs = model_quant.create_config(kwargs, allow_quant, modules_to_not_convert=["distilled_guidance_layer"])
+    if fn.endswith('.safetensors') and os.path.isfile(fn):
+        pipe = diffusers.ChromaPipeline.from_single_file(fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
+    else:
+        pipe = cls.from_pretrained(repo_id or fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
+
+    if shared.opts.teacache_enabled and model_quant.check_nunchaku('Model'):
+        from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
+        apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
+
+    # release memory
+    transformer = None
+    text_encoder = None
+    vae = None
+    for k in kwargs.keys():
+        kwargs[k] = None
+    sd_hijack_te.init_hijack(pipe)
+    devices.torch_gc(force=True)
+    return pipe
diff --git a/modules/model_cogview.py b/modules/model_cogview.py
index 8b761d816..eaba437ed 100644
--- a/modules/model_cogview.py
+++ b/modules/model_cogview.py
@@ -7,7 +7,7 @@ def load_cogview3(checkpoint_info, diffusers_load_config={}):
     modelloader.hf_login()
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
 
-    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
     shared.log.debug(f'Load model: type=CogView3 transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
     transformer = diffusers.CogView3PlusTransformer2DModel.from_pretrained(
         repo_id,
@@ -44,7 +44,7 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}):
     modelloader.hf_login()
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
 
-    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
     shared.log.debug(f'Load model: type=CogView4 transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
     transformer = diffusers.CogView4Transformer2DModel.from_pretrained(
         repo_id,
diff --git a/modules/model_cosmos.py b/modules/model_cosmos.py
new file mode 100644
index 000000000..6d7b17f33
--- /dev/null
+++ b/modules/model_cosmos.py
@@ -0,0 +1,108 @@
+import os
+import transformers
+import diffusers
+from huggingface_hub import auth_check
+from modules import shared, devices, sd_models, model_quant, modelloader, sd_hijack_te
+
+
+def load_transformer(repo_id, diffusers_load_config={}):
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True)
+    fn = None
+
+    if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
+        from modules import sd_unet
+        if shared.opts.sd_unet not in list(sd_unet.unet_dict):
+            shared.log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}')
+            return None
+        fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None
+
+    if fn is not None and 'gguf' in fn.lower():
+        shared.log.error('Load model: type=Cosmos format="gguf" unsupported')
+        transformer = None
+    elif fn is not None and 'safetensors' in fn.lower():
+        shared.log.debug(f'Load model: type=Cosmos transformer="{repo_id}" quant="{model_quant.get_quant(repo_id)}" args={load_args}')
+        transformer = diffusers.CosmosTransformer3DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args)
+    else:
+        shared.log.debug(f'Load model: type=Cosmos transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+        transformer = diffusers.CosmosTransformer3DModel.from_pretrained(
+            repo_id,
+            subfolder="transformer",
+            cache_dir=shared.opts.hfcache_dir,
+            **load_args,
+            **quant_args,
+        )
+    if shared.opts.diffusers_offload_mode != 'none' and transformer is not None:
+        sd_models.move_model(transformer, devices.cpu)
+    return transformer
+
+
+def load_text_encoder(repo_id, diffusers_load_config={}):
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
+    shared.log.debug(f'Load model: type=Cosmos te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+    text_encoder = transformers.T5EncoderModel.from_pretrained(
+        repo_id,
+        subfolder="text_encoder",
+        cache_dir=shared.opts.hfcache_dir,
+        **load_args,
+        **quant_args,
+    )
+    if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None:
+        sd_models.move_model(text_encoder, devices.cpu)
+
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
+    llama_repo = shared.opts.model_h1_llama_repo if shared.opts.model_h1_llama_repo != 'Default' else 'meta-llama/Meta-Llama-3.1-8B-Instruct'
+    shared.log.debug(f'Load model: type=HiDream te4="{llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+
+    return text_encoder
+
+
+def load_cosmos_t2i(checkpoint_info, diffusers_load_config={}):
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    login = modelloader.hf_login()
+    try:
+        auth_check(repo_id)
+    except Exception as e:
+        shared.log.error(f'Load model: repo="{repo_id}" login={login} {e}')
+        return False
+
+    transformer = load_transformer(repo_id, diffusers_load_config)
+    text_encoder = load_text_encoder(repo_id, diffusers_load_config)
+    safety_checker = Fake_safety_checker()
+
+    load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
+    shared.log.debug(f'Load model: type=Cosmos model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+    cls = diffusers.Cosmos2TextToImagePipeline
+    pipe = cls.from_pretrained(
+        repo_id,
+        transformer=transformer,
+        text_encoder=text_encoder,
+        safety_checker=safety_checker,
+        cache_dir=shared.opts.diffusers_dir,
+        **load_args,
+    )
+
+    sd_hijack_te.init_hijack(pipe)
+    del text_encoder
+    del transformer
+
+    devices.torch_gc()
+    return pipe
+
+
+class Fake_safety_checker:
+    def __init__(self):
+        from diffusers.utils import import_utils
+        import_utils._cosmos_guardrail_available = True # pylint: disable=protected-access
+
+    def __call__(self, *args, **kwargs): # pylint: disable=unused-argument
+        return
+
+    def to(self, _device):
+        pass
+
+    def check_text_safety(self, _prompt):
+        return True
+
+    def check_video_safety(self, vid):
+        return vid
diff --git a/modules/model_flex.py b/modules/model_flex.py
index 953a22ec1..1ce5d6145 100644
--- a/modules/model_flex.py
+++ b/modules/model_flex.py
@@ -6,7 +6,7 @@ from modules import shared, devices, sd_models, model_quant, modelloader, sd_hij
 
 
 def load_transformer(repo_id, diffusers_load_config={}):
-    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer', device_map=True)
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True)
     fn = None
 
     if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
@@ -24,7 +24,7 @@ def load_transformer(repo_id, diffusers_load_config={}):
     elif fn is not None and 'safetensors' in fn.lower():
         shared.log.debug(f'Load model: type=FLEX transformer="{repo_id}" quant="{model_quant.get_quant(repo_id)}" args={load_args}')
         transformer = diffusers.FluxTransformer2DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args)
-    # elif model_quant.check_nunchaku('Transformer'):
+    # elif model_quant.check_nunchaku('Model'):
     #     shared.log.error(f'Load model: type=HiDream transformer="{repo_id}" quant="Nunchaku" unsupported')
     #     transformer = None
     else:
diff --git a/modules/model_flux.py b/modules/model_flux.py
index b5f482f2b..518e396d7 100644
--- a/modules/model_flux.py
+++ b/modules/model_flux.py
@@ -110,7 +110,7 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
 
 def load_quants(kwargs, repo_id, cache_dir, allow_quant):
     try:
-        if 'transformer' not in kwargs and model_quant.check_nunchaku('Transformer'):
+        if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'):
             import nunchaku
             nunchaku_precision = nunchaku.utils.get_precision()
             nunchaku_repo = None
@@ -128,8 +128,8 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant):
                 kwargs['transformer'].quantization_method = 'SVDQuant'
                 if shared.opts.nunchaku_attention:
                     kwargs['transformer'].set_attention_impl("nunchaku-fp16")
-        elif 'transformer' not in kwargs and model_quant.check_quant('Transformer'):
-            quant_args = model_quant.create_config(allow=allow_quant, module='Transformer')
+        elif 'transformer' not in kwargs and model_quant.check_quant('Model'):
+            quant_args = model_quant.create_config(allow=allow_quant, module='Model')
             if quant_args:
                 kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
         if 'text_encoder_2' not in kwargs and model_quant.check_nunchaku('TE'):
@@ -138,6 +138,7 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant):
             nunchaku_repo = 'mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors'
             shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}')
             kwargs['text_encoder_2'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
+            kwargs['text_encoder_2'].quantization_method = 'SVDQuant'
         elif 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'):
             quant_args = model_quant.create_config(allow=allow_quant, module='TE')
             if quant_args:
@@ -186,7 +187,7 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
             transformer, _text_encoder_2 = load_flux_nf4(file_path, prequantized=False)
             if transformer is not None:
                 return transformer
-        quant_args = model_quant.create_config(module='Transformer')
+        quant_args = model_quant.create_config(module='Model')
         if quant_args:
             shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=torchao dtype={devices.dtype}')
             transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config, **quant_args)
@@ -317,6 +318,13 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
         cls = diffusers.FluxControlPipeline
     elif 'Depth' in repo_id:
         cls = diffusers.FluxControlPipeline
+    elif 'Kontext' in repo_id:
+        cls = diffusers.FluxKontextPipeline
+        from diffusers import pipelines
+        pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline
+        pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline
+        pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline
+
     else:
         cls = diffusers.FluxPipeline
     shared.log.debug(f'Load model: type=FLUX cls={cls.__name__} preloaded={list(kwargs)} revision={diffusers_load_config.get("revision", None)}')
@@ -340,7 +348,7 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     else:
         pipe = cls.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
 
-    if shared.opts.teacache_enabled and model_quant.check_nunchaku('Transformer'):
+    if shared.opts.teacache_enabled and model_quant.check_nunchaku('Model'):
         from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
         apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
 
diff --git a/modules/model_hidream.py b/modules/model_hidream.py
index f3c5c6483..ac7b335d1 100644
--- a/modules/model_hidream.py
+++ b/modules/model_hidream.py
@@ -6,7 +6,7 @@ from modules import shared, devices, sd_models, model_quant, modelloader, sd_hij
 
 
 def load_transformer(repo_id, diffusers_load_config={}):
-    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer', device_map=True)
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True)
     fn = None
 
     if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
@@ -24,7 +24,7 @@ def load_transformer(repo_id, diffusers_load_config={}):
     elif fn is not None and 'safetensors' in fn.lower():
         shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant(repo_id)}" args={load_args}')
         transformer = diffusers.HiDreamImageTransformer2DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args)
-    # elif model_quant.check_nunchaku('Transformer'):
+    # elif model_quant.check_nunchaku('Model'):
     #     shared.log.error(f'Load model: type=HiDream transformer="{repo_id}" quant="Nunchaku" unsupported')
     #     transformer = None
     else:
@@ -56,7 +56,7 @@ def load_text_encoders(repo_id, diffusers_load_config={}):
     if shared.opts.diffusers_offload_mode != 'none' and text_encoder_3 is not None:
         sd_models.move_model(text_encoder_3, devices.cpu)
 
-    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True)
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
     llama_repo = shared.opts.model_h1_llama_repo if shared.opts.model_h1_llama_repo != 'Default' else 'meta-llama/Meta-Llama-3.1-8B-Instruct'
     shared.log.debug(f'Load model: type=HiDream te4="{llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
 
diff --git a/modules/model_lumina.py b/modules/model_lumina.py
index d817fa48c..306881bbd 100644
--- a/modules/model_lumina.py
+++ b/modules/model_lumina.py
@@ -2,13 +2,12 @@ import os
 import transformers
 import diffusers
 from huggingface_hub import repo_exists
-from modules import errors, shared, sd_unet, sd_hijack_te
+from modules import errors, shared, sd_models, sd_unet, sd_hijack_te, devices, modelloader, model_quant
 
 debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
 def load_lumina(_checkpoint_info, diffusers_load_config={}):
-    from modules import shared, devices, modelloader, model_quant
     modelloader.hf_login()
     load_config, _quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
     pipe = diffusers.LuminaText2ImgPipeline.from_pretrained(
@@ -21,7 +20,6 @@ def load_lumina(_checkpoint_info, diffusers_load_config={}):
 
 
 def load_lumina2(checkpoint_info, diffusers_load_config={}):
-    from modules import shared, devices, sd_models, model_quant
     transformer, text_encoder, vae = None, None, None
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
     if os.path.isdir(checkpoint_info.filename) and not repo_exists(repo_id):
@@ -32,7 +30,7 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
         shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.Lumina2Transformer2DModel.__name__}')
         diffusers.Lumina2Transformer2DModel.forward = teacache.teacache_lumina2_forward # patch must be done before transformer is loaded
 
-    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
+    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model')
     if shared.opts.sd_unet != 'Default':
         try:
             debug(f'Load model: type=Lumina2 unet="{shared.opts.sd_unet}"')
diff --git a/modules/model_omnigen.py b/modules/model_omnigen.py
index b7eb4684e..b2fb95830 100644
--- a/modules/model_omnigen.py
+++ b/modules/model_omnigen.py
@@ -1,25 +1,32 @@
-def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
-    from modules import shared, devices, sd_models, shared_items
-    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+import os
+import diffusers
+from modules import errors, shared, devices, sd_models, model_quant
 
-    # load
-    from modules.omnigen import OmniGenPipeline
-    shared_items.pipelines['OmniGen'] = OmniGenPipeline
-    pipe = OmniGenPipeline.from_pretrained(
-        model_name=repo_id,
-        vae_path='madebyollin/sdxl-vae-fp16-fix',
+debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
+
+
+def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    vae = None
+
+    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model')
+    transformer = diffusers.OmniGenTransformer2DModel.from_pretrained(
+        repo_id,
+        subfolder="transformer",
         cache_dir=shared.opts.diffusers_dir,
+        **load_config,
+        **quant_config,
+    )
+
+    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+    if vae is not None:
+        load_config['vae'] = vae
+    pipe = diffusers.OmniGenPipeline.from_pretrained(
+        repo_id,
+        transformer=transformer,
+        cache_dir=shared.opts.diffusers_dir,
+        **load_config,
     )
 
-    # init
-    pipe.device = devices.device
-    pipe.dtype = devices.dtype
-    pipe.model.device = devices.device
-    pipe.separate_cfg_infer = True
-    pipe.use_kv_cache = False
-    pipe.model.to(device=devices.device, dtype=devices.dtype)
-    if shared.opts.diffusers_eval:
-        pipe.model.eval()
-    pipe.vae.to(devices.device, dtype=devices.dtype)
     devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_omnigen2.py b/modules/model_omnigen2.py
new file mode 100644
index 000000000..285e46f50
--- /dev/null
+++ b/modules/model_omnigen2.py
@@ -0,0 +1,49 @@
+import os
+from modules import shared, devices, sd_models, model_quant
+
+debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
+
+
+def load_omnigen2(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+
+    from modules.omnigen2 import OmniGen2Pipeline, OmniGen2Transformer2DModel, Qwen2_5_VLForConditionalGeneration
+    import diffusers
+    from diffusers import pipelines
+    diffusers.OmniGen2Pipeline = OmniGen2Pipeline # monkey-pathch
+    pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
+    pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
+    pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen2"] = diffusers.OmniGen2Pipeline
+
+    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model')
+    transformer = OmniGen2Transformer2DModel.from_pretrained(
+        repo_id,
+        subfolder="transformer",
+        cache_dir=shared.opts.diffusers_dir,
+        trust_remote_code=True,
+        **load_config,
+        **quant_config,
+    )
+
+    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='TE')
+    mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        repo_id,
+        subfolder="mllm",
+        cache_dir=shared.opts.diffusers_dir,
+        trust_remote_code=True,
+        **load_config,
+        **quant_config,
+    )
+
+    pipe = OmniGen2Pipeline.from_pretrained(
+        repo_id,
+        # transformer=transformer,
+        mllm=mllm,
+        cache_dir=shared.opts.diffusers_dir,
+        trust_remote_code=True,
+        **load_config,
+    )
+    pipe.transformer = transformer # for omnigen2 transformer must be loaded after pipeline
+
+    devices.torch_gc(force=True)
+    return pipe
diff --git a/modules/model_pixart.py b/modules/model_pixart.py
index 6f6d6cf1c..7edcee7ac 100644
--- a/modules/model_pixart.py
+++ b/modules/model_pixart.py
@@ -15,7 +15,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
     if not file_exists(repo_id_pipe, "model_index.json"):
         repo_id_pipe = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
 
-    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
+    load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
     transformer = diffusers.PixArtTransformer2DModel.from_pretrained(
         repo_id,
         subfolder='transformer',
diff --git a/modules/model_quant.py b/modules/model_quant.py
index 3c57ac67c..cbacd7842 100644
--- a/modules/model_quant.py
+++ b/modules/model_quant.py
@@ -37,124 +37,122 @@ def get_quant(name):
     return 'none'
 
 
-def create_bnb_config(kwargs = None, allow_bnb: bool = True, module: str = 'Model'):
+def create_bnb_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
     from modules import shared, devices
-    if len(shared.opts.bnb_quantization) > 0 and allow_bnb:
-        if 'Model' in shared.opts.bnb_quantization or (module is not None and module in shared.opts.bnb_quantization) or module == 'any':
-            load_bnb()
-            if bnb is None:
-                return kwargs
-            bnb_config = diffusers.BitsAndBytesConfig(
-                load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
-                load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
-                bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
-                bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
-                bnb_4bit_compute_dtype=devices.dtype
-            )
-            log.debug(f'Quantization: module={module} type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-            if kwargs is None:
-                return bnb_config
-            else:
-                kwargs['quantization_config'] = bnb_config
-                return kwargs
+    if allow and (module == 'any' or module in shared.opts.bnb_quantization):
+        load_bnb()
+        if bnb is None:
+            return kwargs
+        bnb_config = diffusers.BitsAndBytesConfig(
+            load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
+            load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
+            bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
+            bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
+            bnb_4bit_compute_dtype=devices.dtype,
+            llm_int8_skip_modules=modules_to_not_convert,
+        )
+        log.debug(f'Quantization: module={module} type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
+        if kwargs is None:
+            return bnb_config
+        else:
+            kwargs['quantization_config'] = bnb_config
+            return kwargs
     return kwargs
 
 
-def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model'):
+def create_ao_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
     from modules import shared
-    if len(shared.opts.torchao_quantization) > 0 and (shared.opts.torchao_quantization_mode == 'pre') and allow_ao:
-        if 'Model' in shared.opts.torchao_quantization or (module is not None and module in shared.opts.torchao_quantization) or module == 'any':
-            torchao = load_torchao()
-            if torchao is None:
-                return kwargs
-            if module in {'TE', 'LLM'}:
-                ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type)
-            else:
-                ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type)
-            log.debug(f'Quantization: module={module} type=torchao dtype={shared.opts.torchao_quantization_type}')
-            if kwargs is None:
-                return ao_config
-            else:
-                kwargs['quantization_config'] = ao_config
-                return kwargs
+    if allow and (shared.opts.torchao_quantization_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.torchao_quantization):
+        torchao = load_torchao()
+        if torchao is None:
+            return kwargs
+        if module in {'TE', 'LLM'}:
+            ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
+        else:
+            ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
+        log.debug(f'Quantization: module={module} type=torchao dtype={shared.opts.torchao_quantization_type}')
+        if kwargs is None:
+            return ao_config
+        else:
+            kwargs['quantization_config'] = ao_config
+            return kwargs
     return kwargs
 
 
-def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str = 'Model'):
+def create_quanto_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
     from modules import shared
-    if len(shared.opts.quanto_quantization) > 0 and allow_quanto:
-        if 'Model' in shared.opts.quanto_quantization or (module is not None and module in shared.opts.quanto_quantization) or module == 'any':
-            load_quanto(silent=True)
-            if optimum_quanto is None:
-                return kwargs
-            if module in {'TE', 'LLM'}:
-                quanto_config = transformers.QuantoConfig(weights=shared.opts.quanto_quantization_type)
-                quanto_config.weights_dtype = quanto_config.weights
-            else:
-                quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type)
-                quanto_config.activations = None # patch so it works with transformers
-                quanto_config.weights = quanto_config.weights_dtype
-            log.debug(f'Quantization: module={module} type=quanto dtype={shared.opts.quanto_quantization_type}')
-            if kwargs is None:
-                return quanto_config
-            else:
-                kwargs['quantization_config'] = quanto_config
-                return kwargs
+    if allow and (module == 'any' or module in shared.opts.quanto_quantization):
+        load_quanto(silent=True)
+        if optimum_quanto is None:
+            return kwargs
+        if module in {'TE', 'LLM'}:
+            quanto_config = transformers.QuantoConfig(weights=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
+            quanto_config.weights_dtype = quanto_config.weights
+        else:
+            quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
+            quanto_config.activations = None # patch so it works with transformers
+            quanto_config.weights = quanto_config.weights_dtype
+        log.debug(f'Quantization: module={module} type=quanto dtype={shared.opts.quanto_quantization_type}')
+        if kwargs is None:
+            return quanto_config
+        else:
+            kwargs['quantization_config'] = quanto_config
+            return kwargs
     return kwargs
 
 
-def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Model', weights_dtype: str = None):
+def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = []):
     from modules import devices, shared
-    if len(shared.opts.sdnq_quantize_weights) > 0 and (shared.opts.sdnq_quantize_mode == 'pre') and allow_sdnq:
-        if 'Model' in shared.opts.sdnq_quantize_weights or (module is not None and module in shared.opts.sdnq_quantize_weights) or module == 'any':
-            from modules.sdnq import SDNQQuantizer, SDNQConfig
-            diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
-            transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
-            diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
-            transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
+    if allow and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.sdnq_quantize_weights):
+        from modules.sdnq import SDNQQuantizer, SDNQConfig
+        diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
+        transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
+        diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
+        transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
 
-            if weights_dtype is None:
-                if shared.opts.sdnq_quantize_weights_mode_te != "default" and module in {"TE", "LLM"}:
-                    weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
-                else:
-                    weights_dtype = shared.opts.sdnq_quantize_weights_mode
-
-            if shared.opts.device_map == "gpu":
-                quantization_device = devices.device
-                return_device = devices.device
-            elif shared.opts.diffusers_offload_mode in {"none", "model"}:
-                quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
-                return_device = devices.device
-            elif shared.opts.sdnq_quantize_with_gpu:
-                quantization_device = devices.device
-                return_device = devices.cpu
+        if weights_dtype is None:
+            if module in {"TE", "LLM"} and shared.opts.sdnq_quantize_weights_mode_te not in {"same as model", "default"}:
+                weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
             else:
-                quantization_device = None
-                return_device = None
+                weights_dtype = shared.opts.sdnq_quantize_weights_mode
+        if weights_dtype is None or weights_dtype == 'none':
+            return kwargs
 
-            sdnq_config = SDNQConfig(
-                weights_dtype=weights_dtype,
-                group_size=shared.opts.sdnq_quantize_weights_group_size,
-                quant_conv=shared.opts.sdnq_quantize_conv_layers,
-                use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
-                use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
-                dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
-                quantization_device=quantization_device,
-                return_device=return_device,
-            )
-            log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype}')
-            if kwargs is None:
-                return sdnq_config
-            else:
-                kwargs['quantization_config'] = sdnq_config
-                return kwargs
+        if shared.opts.device_map == "gpu":
+            quantization_device = devices.device
+            return_device = devices.device
+        elif shared.opts.diffusers_offload_mode in {"none", "model"}:
+            quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
+            return_device = devices.device
+        elif shared.opts.sdnq_quantize_with_gpu:
+            quantization_device = devices.device
+            return_device = devices.cpu
+        else:
+            quantization_device = None
+            return_device = None
+
+        sdnq_config = SDNQConfig(
+            weights_dtype=weights_dtype,
+            group_size=shared.opts.sdnq_quantize_weights_group_size,
+            quant_conv=shared.opts.sdnq_quantize_conv_layers,
+            use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
+            use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
+            dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
+            quantization_device=quantization_device,
+            return_device=return_device,
+            modules_to_not_convert=modules_to_not_convert,
+        )
+        log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}')
+        if kwargs is None:
+            return sdnq_config
+        else:
+            kwargs['quantization_config'] = sdnq_config
+            return kwargs
     return kwargs
 
 
 def check_quant(module: str = ''):
     from modules import shared
-    if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization or 'Model' in shared.opts.sdnq_quantize_weights:
-        return True
     if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization or module in shared.opts.sdnq_quantize_weights:
         return True
     return False
@@ -162,7 +160,7 @@ def check_quant(module: str = ''):
 
 def check_nunchaku(module: str = ''):
     from modules import shared
-    if 'Model' not in shared.opts.nunchaku_quantization and module not in shared.opts.nunchaku_quantization:
+    if module not in shared.opts.nunchaku_quantization:
         return False
     from modules import mit_nunchaku
     mit_nunchaku.install_nunchaku()
@@ -171,25 +169,25 @@ def check_nunchaku(module: str = ''):
     return True
 
 
-def create_config(kwargs = None, allow: bool = True, module: str = 'Model'):
+def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert = []):
     if kwargs is None:
         kwargs = {}
-    kwargs = create_sdnq_config(kwargs, allow_sdnq=allow, module=module)
+    kwargs = create_sdnq_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
     if kwargs is not None and 'quantization_config' in kwargs:
         if debug:
             log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}')
         return kwargs
-    kwargs = create_bnb_config(kwargs, allow_bnb=allow, module=module)
+    kwargs = create_bnb_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
     if kwargs is not None and 'quantization_config' in kwargs:
         if debug:
             log.trace(f'Quantization: type=bnb config={kwargs.get("quantization_config", None)}')
         return kwargs
-    kwargs = create_quanto_config(kwargs, allow_quanto=allow, module=module)
+    kwargs = create_quanto_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
     if kwargs is not None and 'quantization_config' in kwargs:
         if debug:
             log.trace(f'Quantization: type=quanto config={kwargs.get("quantization_config", None)}')
         return kwargs
-    kwargs = create_ao_config(kwargs, allow_ao=allow, module=module)
+    kwargs = create_ao_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
     if kwargs is not None and 'quantization_config' in kwargs:
         if debug:
             log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}')
@@ -305,13 +303,13 @@ def apply_layerwise(sd_model, quiet:bool=False):
                     m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
                     m.quantization_method = 'LayerWise'
                     log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
-            if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization or 'Transformer' in shared.opts.layerwise_quantization):
+            if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization):
                 m = getattr(sd_model, module)
                 if hasattr(m, 'enable_layerwise_casting'):
                     m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
                     m.quantization_method = 'LayerWise'
                     log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
-            if module.startswith('text_encoder') and ('Model' in shared.opts.layerwise_quantization or 'TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
+            if module.startswith('text_encoder') and ('TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
                 m = getattr(sd_model, module)
                 if hasattr(m, 'enable_layerwise_casting'):
                     m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
@@ -322,30 +320,21 @@ def apply_layerwise(sd_model, quiet:bool=False):
                 log.error(f'Quantization: type=layerwise {e}')
 
 
-def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
+def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weights_dtype: str = None, modules_to_not_convert: list = []):
     global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
     from modules import devices, shared
     from modules.sdnq import apply_sdnq_to_module
 
-    model.eval()
+    if weights_dtype is None:
+        if op is not None and ("text_encoder" in op or op in {"TE", "LLM"}) and shared.opts.sdnq_quantize_weights_mode_te not in {"same as model", "default"}:
+            weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
+        else:
+            weights_dtype = shared.opts.sdnq_quantize_weights_mode
 
-    if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
-        import torch
-        from modules.sdnq import SDNQ_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
-        for i in range(len(model.encoder.block)):
-            model.encoder.block[i].layer[1].DenseReluDense = SDNQ_T5DenseGatedActDense(
-                model.encoder.block[i].layer[1].DenseReluDense,
-                dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
-            )
-
-    backup_embeddings = None
-    if hasattr(model, "get_input_embeddings"):
-        backup_embeddings = copy.deepcopy(model.get_input_embeddings())
-
-    if shared.opts.sdnq_quantize_weights_mode_te != "default" and op is not None and "text_encoder" in op:
-        weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
-    else:
-        weights_dtype = shared.opts.sdnq_quantize_weights_mode
+    if weights_dtype is None or weights_dtype == 'none':
+        return model
+    if debug:
+        log.trace(f'Quantization: type=SDNQ op={op} cls={model.__class__} dtype={weights_dtype} mode{shared.opts.diffusers_offload_mode}')
 
     if shared.opts.diffusers_offload_mode in {"none", "model"}:
         quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
@@ -357,6 +346,16 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
         quantization_device = None
         return_device = None
 
+    if getattr(model, "_keep_in_fp32_modules", None) is not None:
+        modules_to_not_convert.extend(model._keep_in_fp32_modules)
+    if model.__class__.__name__ == "ChromaTransformer2DModel":
+        modules_to_not_convert.append("distilled_guidance_layer")
+
+    model.eval()
+    backup_embeddings = None
+    if hasattr(model, "get_input_embeddings"):
+        backup_embeddings = copy.deepcopy(model.get_input_embeddings())
+
     model = apply_sdnq_to_module(
         model,
         weights_dtype=weights_dtype,
@@ -369,6 +368,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
         quantization_device=quantization_device,
         return_device=return_device,
         param_name=op,
+        modules_to_not_convert=modules_to_not_convert,
     )
     model.quantization_method = 'SDNQ'
 
@@ -402,7 +402,7 @@ def sdnq_quantize_weights(sd_model):
     try:
         t0 = time.time()
         from modules import shared, devices, sd_models
-        log.info(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights}")
+        log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul}  group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32}")
         global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
 
         sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq")
@@ -427,8 +427,12 @@ def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activation
     from modules import devices, shared
     quanto = load_quanto('Quantize model: type=Optimum Quanto')
     global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
-    if sd_model is not None and "Flux" in sd_model.__class__.__name__: # LayerNorm is not supported
+    if sd_model is not None and ("Flux" in sd_model.__class__.__name__ or "Chroma" in sd_model.__class__.__name__): # LayerNorm is not supported
         exclude_list = ["transformer_blocks.*.norm1.norm", "transformer_blocks.*.norm2", "transformer_blocks.*.norm1_context.norm", "transformer_blocks.*.norm2_context", "single_transformer_blocks.*.norm.norm", "norm_out.norm"]
+        if "Chroma" in sd_model.__class__.__name__:
+            # we ignore the distilled guidance layer because it degrades quality too much
+            # see: https://github.com/huggingface/diffusers/pull/11698#issuecomment-2969717180 for more details
+            exclude_list.append("distilled_guidance_layer.*")
     else:
         exclude_list = None
     weights = getattr(quanto, weights) if weights is not None else getattr(quanto, shared.opts.optimum_quanto_weights_type)
@@ -549,7 +553,7 @@ def torchao_quantization(sd_model):
     return sd_model
 
 
-def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, allow_quant:bool=True):
+def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, allow_quant:bool=True, modules_to_not_convert: list = []):
     from modules import shared, devices
     config = load_config.copy()
     if 'torch_dtype' not in config:
@@ -565,26 +569,26 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al
     # if 'variant' in config:
     #     del config['variant']
     if device_map:
-        if devices.backend == "ipex" and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}:
+        if devices.backend == 'ipex' and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}:
             config['device_map'] = 'cpu' # alchemist gpus hits the 4GB allocation limit with transformers, UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS emulates above 4GB allocations
         elif shared.opts.device_map == 'cpu':
             config['device_map'] = 'cpu'
         elif shared.opts.device_map == 'gpu':
             config['device_map'] = devices.device
     if allow_quant:
-        quant_args = create_config(module=module)
+        quant_args = create_config(module=module, modules_to_not_convert=modules_to_not_convert)
     else:
         quant_args = {}
     return config, quant_args
 
 
-def do_post_load_quant(sd_model):
+def do_post_load_quant(sd_model, allow=True):
     from modules import shared
-    if shared.opts.sdnq_quantize_weights and shared.opts.sdnq_quantize_mode == 'post':
+    if shared.opts.sdnq_quantize_weights and (shared.opts.sdnq_quantize_mode == 'post' or (allow and shared.opts.sdnq_quantize_mode == 'auto')):
         sd_model = sdnq_quantize_weights(sd_model)
     if shared.opts.optimum_quanto_weights:
         sd_model = optimum_quanto_weights(sd_model)
-    if shared.opts.torchao_quantization and shared.opts.torchao_quantization_mode == 'post':
+    if shared.opts.torchao_quantization and (shared.opts.torchao_quantization_mode == 'post' or (allow and shared.opts.torchao_quantization_mode == 'auto')):
         sd_model = torchao_quantization(sd_model)
     if shared.opts.layerwise_quantization:
         apply_layerwise(sd_model)
diff --git a/modules/model_sana.py b/modules/model_sana.py
index d211321fd..4de29f0d8 100644
--- a/modules/model_sana.py
+++ b/modules/model_sana.py
@@ -7,14 +7,14 @@ from modules import shared, sd_models, devices, modelloader, model_quant
 
 def load_quants(kwargs, repo_id, cache_dir):
     kwargs_copy = kwargs.copy()
-    if model_quant.check_nunchaku('Transformer') and 'Sana_1600M' in repo_id: # only sana-1600m
+    if model_quant.check_nunchaku('Model') and 'Sana_1600M' in repo_id: # only sana-1600m
         import nunchaku
         nunchaku_precision = nunchaku.utils.get_precision()
         nunchaku_repo = f"mit-han-lab/svdq-{nunchaku_precision}-sana-1600m"
         shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} attention={shared.opts.nunchaku_attention}')
         kwargs['transformer'] = nunchaku.NunchakuSanaTransformer2DModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
-    elif model_quant.check_quant('Transformer'):
-        load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='Transformer')
+    elif model_quant.check_quant('Model'):
+        load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='Model')
         if quant_args:
             kwargs['transformer'] = diffusers.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args)
     if model_quant.check_quant('TE'):
diff --git a/modules/model_sd3.py b/modules/model_sd3.py
index ecfc533a1..8bfea1afa 100644
--- a/modules/model_sd3.py
+++ b/modules/model_sd3.py
@@ -57,7 +57,7 @@ def load_overrides(kwargs, cache_dir):
 
 
 def load_quants(kwargs, repo_id, cache_dir):
-    quant_args = model_quant.create_config(module='Transformer')
+    quant_args = model_quant.create_config(module='Model')
     if quant_args and 'quantization_config' in quant_args:
         kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
     quant_args = model_quant.create_config(module='TE')
diff --git a/modules/model_te.py b/modules/model_te.py
index 6472721ee..27ad10357 100644
--- a/modules/model_te.py
+++ b/modules/model_te.py
@@ -4,7 +4,6 @@ import torch
 import transformers
 from safetensors.torch import load_file
 from modules import shared, devices, files_cache, errors, model_quant
-from installer import install
 
 
 te_dict = {}
@@ -72,27 +71,32 @@ def load_t5(name=None, cache_dir=None):
     elif 'int8' in name.lower():
         from modules.model_quant import create_sdnq_config
         quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='int8')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif 'uint4' in name.lower():
         from modules.model_quant import create_sdnq_config
         quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='uint4')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif 'qint4' in name.lower():
         model_quant.load_quanto('Load model: type=T5')
         quantization_config = transformers.QuantoConfig(weights='int4')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif 'qint8' in name.lower():
         model_quant.load_quanto('Load model: type=T5')
         quantization_config = transformers.QuantoConfig(weights='int8')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif '/' in name:
         shared.log.debug(f'Load model: type=T5 repo={name}')
         quant_config = model_quant.create_config(module='TE')
-        t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config)
 
     else:
         t5 = None
diff --git a/modules/modeldata.py b/modules/modeldata.py
index 688cff23d..789cccfef 100644
--- a/modules/modeldata.py
+++ b/modules/modeldata.py
@@ -27,12 +27,16 @@ def get_model_type(pipe):
         model_type = 'sc'
     elif "AuraFlow" in name:
         model_type = 'auraflow'
+    elif 'Chroma' in name:
+        model_type = 'chroma'
     elif "Flux" in name or "Flex1" in name or "Flex2" in name:
         model_type = 'f1'
     elif "Lumina2" in name:
         model_type = 'lumina2'
     elif "Lumina" in name:
         model_type = 'lumina'
+    elif "OmniGen2" in name:
+        model_type = 'omnigen2'
     elif "OmniGen" in name:
         model_type = 'omnigen'
     elif "CogView3" in name:
@@ -43,6 +47,8 @@ def get_model_type(pipe):
         model_type = 'sana'
     elif "HiDream" in name:
         model_type = 'h1'
+    elif "Cosmos2TextToImage" in name:
+        model_type = 'cosmos'
     # video models
     elif "CogVideo" in name:
         model_type = 'cogvideo'
diff --git a/modules/omnigen/__init__.py b/modules/omnigen/__init__.py
deleted file mode 100644
index 40315a6f3..000000000
--- a/modules/omnigen/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from .model import OmniGen
-from .processor import OmniGenProcessor
-from .scheduler import OmniGenScheduler
-from .pipeline import OmniGenPipeline
diff --git a/modules/omnigen/model.py b/modules/omnigen/model.py
deleted file mode 100644
index 17d696b53..000000000
--- a/modules/omnigen/model.py
+++ /dev/null
@@ -1,390 +0,0 @@
-# The code is revised from DiT
-import os
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-from safetensors.torch import load_file
-from diffusers.loaders import PeftAdapterMixin
-from huggingface_hub import snapshot_download
-from .transformer import Phi3Config, Phi3Transformer
-
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-class TimestepEmbedder(nn.Module):
-    """
-    Embeds scalar timesteps into vector representations.
-    """
-    def __init__(self, hidden_size, frequency_embedding_size=256):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=True),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        """
-        Create sinusoidal timestep embeddings.
-        :param t: a 1-D Tensor of N indices, one per batch element.
-                          These may be fractional.
-        :param dim: the dimension of the output.
-        :param max_period: controls the minimum frequency of the embeddings.
-        :return: an (N, D) Tensor of positional embeddings.
-        """
-        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
-        ).to(device=t.device)
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t, dtype=torch.float32):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-
-class FinalLayer(nn.Module):
-    """
-    The final layer of DiT.
-    """
-    def __init__(self, hidden_size, patch_size, out_channels):
-        super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-        self.adaLN_modulation = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
-        )
-
-    def forward(self, x, c):
-        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
-        x = modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
-
-
-def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=1):
-    """
-    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
-    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
-    """
-    if isinstance(grid_size, int):
-        grid_size = (grid_size, grid_size)
-
-    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale
-    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-
-    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
-    return pos_embed
-
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    assert embed_dim % 2 == 0
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
-
-    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
-    return emb
-
-
-def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
-    """
-    embed_dim: output dimension for each position
-    pos: a list of positions to be encoded: size (M,)
-    out: (M, D)
-    """
-    assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float64)
-    omega /= embed_dim / 2.
-    omega = 1. / 10000**omega  # (D/2,)
-
-    pos = pos.reshape(-1)  # (M,)
-    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
-
-    emb_sin = np.sin(out) # (M, D/2)
-    emb_cos = np.cos(out) # (M, D/2)
-
-    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
-    return emb
-
-
-class PatchEmbedMR(nn.Module):
-    """ 2D Image to Patch Embedding
-    """
-    def __init__(
-            self,
-            patch_size: int = 2,
-            in_chans: int = 4,
-            embed_dim: int = 768,
-            bias: bool = True,
-    ):
-        super().__init__()
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
-
-    def forward(self, x):
-        x = self.proj(x)
-        x = x.flatten(2).transpose(1, 2)  # NCHW -> NLC
-        return x
-
-
-class OmniGen(nn.Module, PeftAdapterMixin):
-    """
-    Diffusion model with a Transformer backbone.
-    """
-    def __init__(
-        self,
-        transformer_config: Phi3Config,
-        patch_size=2,
-        in_channels=4,
-        pe_interpolation: float = 1.0,
-        pos_embed_max_size: int = 192,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.patch_size = patch_size
-        self.pos_embed_max_size = pos_embed_max_size
-        hidden_size = transformer_config.hidden_size
-        self.x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
-        self.input_x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
-        self.time_token = TimestepEmbedder(hidden_size)
-        self.t_embedder = TimestepEmbedder(hidden_size)
-        self.pe_interpolation = pe_interpolation
-        pos_embed = get_2d_sincos_pos_embed(hidden_size, pos_embed_max_size, interpolation_scale=self.pe_interpolation, base_size=64)
-        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=True)
-        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
-        self.initialize_weights()
-        self.llm = Phi3Transformer(config=transformer_config)
-        self.llm.config.use_cache = False
-
-    @classmethod
-    def from_pretrained(cls, model_name: str, cache_dir: str=None):
-        if not os.path.exists(os.path.join(model_name, 'model.pt')) and not os.path.exists(os.path.join(model_name, 'model.safetensors')):
-            cache_dir = cache_dir or os.getenv('HF_HUB_CACHE')
-            model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_dir,
-                                           ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
-        config = Phi3Config.from_pretrained(model_name)
-        model = cls(config)
-        if os.path.exists(os.path.join(model_name, 'model.pt')):
-            state_dict = torch.load(os.path.join(model_name, 'model.pt'), map_location='cpu')
-        elif os.path.exists(os.path.join(model_name, 'model.safetensors')):
-            state_dict = load_file(os.path.join(model_name, 'model.safetensors'))
-        else:
-            raise ValueError(f"OmniGen: Could not find model file in {model_name}")
-        model.load_state_dict(state_dict)
-        return model
-
-    def initialize_weights(self):
-        assert not hasattr(self, "llama")
-
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-
-        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
-        w = self.x_embedder.proj.weight.data
-        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
-        nn.init.constant_(self.x_embedder.proj.bias, 0)
-
-        w = self.input_x_embedder.proj.weight.data
-        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
-        nn.init.constant_(self.x_embedder.proj.bias, 0)
-
-
-        # Initialize timestep embedding MLP:
-        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
-        nn.init.normal_(self.time_token.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.time_token.mlp[2].weight, std=0.02)
-
-        # Zero-out output layers:
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
-        nn.init.constant_(self.final_layer.linear.weight, 0)
-        nn.init.constant_(self.final_layer.linear.bias, 0)
-
-    def unpatchify(self, x, h, w):
-        """
-        x: (N, T, patch_size**2 * C)
-        imgs: (N, H, W, C)
-        """
-        c = self.out_channels
-
-        x = x.reshape(shape=(x.shape[0], h//self.patch_size, w//self.patch_size, self.patch_size, self.patch_size, c))
-        x = torch.einsum('nhwpqc->nchpwq', x)
-        imgs = x.reshape(shape=(x.shape[0], c, h, w))
-        return imgs
-
-
-    def cropped_pos_embed(self, height, width):
-        """Crops positional embeddings for SD3 compatibility."""
-        if self.pos_embed_max_size is None:
-            raise ValueError("`pos_embed_max_size` must be set for cropping.")
-
-        height = height // self.patch_size
-        width = width // self.patch_size
-        if height > self.pos_embed_max_size:
-            raise ValueError(
-                f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
-            )
-        if width > self.pos_embed_max_size:
-            raise ValueError(
-                f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
-            )
-
-        top = (self.pos_embed_max_size - height) // 2
-        left = (self.pos_embed_max_size - width) // 2
-        spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)
-        spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]
-        spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
-        return spatial_pos_embed
-
-
-    def patch_multiple_resolutions(self, latents, padding_latent=None, is_input_images:bool=False):
-        if isinstance(latents, list):
-            return_list = False
-            if padding_latent is None:
-                padding_latent = [None] * len(latents)
-                return_list = True
-            patched_latents, num_tokens, shapes = [], [], []
-            for latent, padding in zip(latents, padding_latent):
-                height, width = latent.shape[-2:]
-                if is_input_images:
-                    latent = self.input_x_embedder(latent)
-                else:
-                    latent = self.x_embedder(latent)
-                pos_embed = self.cropped_pos_embed(height, width)
-                latent = latent + pos_embed
-                if padding is not None:
-                    latent = torch.cat([latent, padding], dim=-2)
-                patched_latents.append(latent)
-
-                num_tokens.append(pos_embed.size(1))
-                shapes.append([height, width])
-            if not return_list:
-                latents = torch.cat(patched_latents, dim=0)
-            else:
-                latents = patched_latents
-        else:
-            height, width = latents.shape[-2:]
-            if is_input_images:
-                latents = self.input_x_embedder(latents)
-            else:
-                latents = self.x_embedder(latents)
-            pos_embed = self.cropped_pos_embed(height, width)
-            latents = latents + pos_embed
-            num_tokens = latents.size(1)
-            shapes = [height, width]
-        return latents, num_tokens, shapes
-
-    def forward(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, padding_latent=None, past_key_values=None, return_past_key_values=True):
-        input_is_list = isinstance(x, list)
-        x, num_tokens, shapes = self.patch_multiple_resolutions(x, padding_latent)
-        time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1)
-        if input_img_latents is not None:
-            input_latents, _, _ = self.patch_multiple_resolutions(input_img_latents, is_input_images=True)
-        if input_ids is not None:
-            condition_embeds = self.llm.embed_tokens(input_ids).clone()
-            input_img_inx = 0
-            for b_inx in input_image_sizes.keys():
-                for start_inx, end_inx in input_image_sizes[b_inx]:
-                    condition_embeds[b_inx, start_inx: end_inx] = input_latents[input_img_inx]
-                    input_img_inx += 1
-            if input_img_latents is not None:
-                assert input_img_inx == len(input_latents)
-
-            input_emb = torch.cat([condition_embeds, time_token, x], dim=1)
-        else:
-            input_emb = torch.cat([time_token, x], dim=1)
-        output = self.llm(inputs_embeds=input_emb, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values)
-        output, past_key_values = output.last_hidden_state, output.past_key_values
-        if input_is_list:
-            image_embedding = output[:, -max(num_tokens):]
-            time_emb = self.t_embedder(timestep, dtype=x.dtype)
-            x = self.final_layer(image_embedding, time_emb)
-            latents = []
-            for i in range(x.size(0)):
-                latent = x[i:i+1, :num_tokens[i]]
-                latent = self.unpatchify(latent, shapes[i][0], shapes[i][1])
-                latents.append(latent)
-        else:
-            image_embedding = output[:, -num_tokens:]
-            time_emb = self.t_embedder(timestep, dtype=x.dtype)
-            x = self.final_layer(image_embedding, time_emb)
-            latents = self.unpatchify(x, shapes[0], shapes[1])
-
-        if return_past_key_values:
-            return latents, past_key_values
-        return latents
-
-    @torch.no_grad()
-    def forward_with_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache):
-        """
-        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
-        """
-        self.llm.config.use_cache = use_kv_cache
-        model_out, past_key_values = self.forward(x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, past_key_values=past_key_values, return_past_key_values=True)
-        if use_img_cfg:
-            cond, uncond, img_cond = torch.split(model_out, len(model_out) // 3, dim=0)
-            cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond)
-            model_out = [cond, cond, cond]
-        else:
-            cond, uncond = torch.split(model_out, len(model_out) // 2, dim=0)
-            cond = uncond + cfg_scale * (cond - uncond)
-            model_out = [cond, cond]
-        return torch.cat(model_out, dim=0), past_key_values
-
-
-    @torch.no_grad()
-    def forward_with_separate_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache, return_past_key_values=True):
-        """
-        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
-        """
-        self.llm.config.use_cache = use_kv_cache
-        if past_key_values is None:
-            past_key_values = [None] * len(attention_mask)
-
-        x = torch.split(x, len(x) // len(attention_mask), dim=0)
-        timestep = timestep.to(x[0].dtype)
-        timestep = torch.split(timestep, len(timestep) // len(input_ids), dim=0)
-
-        model_out, pask_key_values = [], []
-        for i in range(len(input_ids)):
-            temp_out, temp_pask_key_values = self.forward(x[i], timestep[i], input_ids[i], input_img_latents[i], input_image_sizes[i], attention_mask[i], position_ids[i], past_key_values[i])
-            model_out.append(temp_out)
-            pask_key_values.append(temp_pask_key_values)
-
-        if len(model_out) == 3:
-            cond, uncond, img_cond = model_out
-            cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond)
-            model_out = [cond, cond, cond]
-        elif len(model_out) == 2:
-            cond, uncond = model_out
-            cond = uncond + cfg_scale * (cond - uncond)
-            model_out = [cond, cond]
-        else:
-            return model_out[0]
-        return torch.cat(model_out, dim=0), pask_key_values
diff --git a/modules/omnigen/pipeline.py b/modules/omnigen/pipeline.py
deleted file mode 100644
index a07467543..000000000
--- a/modules/omnigen/pipeline.py
+++ /dev/null
@@ -1,219 +0,0 @@
-import os
-from typing import List, Union
-from PIL import Image
-import torch
-from huggingface_hub import snapshot_download
-from peft import PeftModel
-from diffusers.models import AutoencoderKL
-from diffusers.utils import replace_example_docstring
-from .model import OmniGen
-from .processor import OmniGenProcessor
-from .scheduler import OmniGenScheduler
-
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> from OmniGen import OmniGenPipeline
-        >>> pipe = FluxControlNetPipeline.from_pretrained(
-        ...     base_model
-        ... )
-        >>> prompt = "A woman holds a bouquet of flowers and faces the camera"
-        >>> image = pipe(
-        ...     prompt,
-        ...     guidance_scale=3.0,
-        ...     num_inference_steps=50,
-        ... ).images[0]
-        >>> image.save("t2i.png")
-        ```
-"""
-
-
-class OmniGenPipeline():
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        model: OmniGen,
-        processor: OmniGenProcessor,
-    ):
-        super().__init__()
-        self.vae = vae
-        self.model = model
-        self.processor = processor
-        self.device = None
-        self.dtype: None
-        self.separate_cfg_infer: bool = True
-        self.use_kv_cache: bool = False
-        # omnigen does not inherit from diffusionpipeline so we hack it
-        self._internal_dict = { # pylint: disable=protected-access
-            'vae': self.vae,
-            'model': self.model,
-            'processor': self.processor,
-        }
-
-    @classmethod
-    def from_pretrained(cls, model_name, vae_path: str=None, cache_dir: str=None):
-        if not os.path.exists(model_name):
-            cache_dir = cache_dir or os.getenv('HF_HUB_CACHE')
-            model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_dir,
-                                           ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
-        model = OmniGen.from_pretrained(model_name)
-        processor = OmniGenProcessor.from_pretrained(model_name)
-        if os.path.exists(os.path.join(model_name, "vae")):
-            vae = AutoencoderKL.from_pretrained(os.path.join(model_name, "vae"))
-        else:
-            vae = AutoencoderKL.from_pretrained(vae_path or "stabilityai/sdxl-vae")
-        return cls(vae, model, processor)
-
-    def merge_lora(self, lora_path: str):
-        model = PeftModel.from_pretrained(self.model, lora_path)
-        model.merge_and_unload()
-        self.model = model
-
-    def to(self, device: Union[str, torch.device]):
-        if isinstance(device, str):
-            device = torch.device(device)
-        self.model.to(device)
-        self.vae.to(device)
-
-    def vae_encode(self, x, dtype):
-        x = x.to(dtype)
-        if self.vae.config.shift_factor is not None:
-            x = self.vae.encode(x).latent_dist.sample()
-            x = (x - self.vae.config.shift_factor) * self.vae.config.scaling_factor
-        else:
-            x = self.vae.encode(x).latent_dist.sample().mul_(self.vae.config.scaling_factor)
-        x = x.to(dtype)
-        return x
-
-    def move_to_device(self, data):
-        if isinstance(data, list):
-            return [x.to(self.device) for x in data]
-        return data.to(self.device)
-
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]],
-        input_images: Union[List[str], List[List[str]]] = None,
-        height: int = 1024,
-        width: int = 1024,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 3,
-        use_img_guidance: bool = True,
-        img_guidance_scale: float = 1.6,
-        output_type: str = 'latent',
-        seed: int = None,
-        ):
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`):
-                The prompt or prompts to guide the image generation. 
-            input_images (`List[str]` or `List[List[str]]`, *optional*):
-                The list of input images. We will replace the "<|image_i|>" in prompt with the 1-th image in list.
-            height (`int`, *optional*, defaults to 1024):
-                The height in pixels of the generated image. The number must be a multiple of 16.
-            width (`int`, *optional*, defaults to 1024):
-                The width in pixels of the generated image. The number must be a multiple of 16.
-            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.
-            guidance_scale (`float`, *optional*, defaults to 4.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.
-            use_img_guidance (`bool`, *optional*, defaults to True):
-                Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800). 
-            img_guidance_scale (`float`, *optional*, defaults to 1.6):
-                Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800). 
-            self.separate_cfg_infer (`bool`, *optional*, defaults to False):
-                Perform inference on images with different guidance separately; this can save memory when generating images of large size at the expense of slower inference.
-            self.use_kv_cache (`bool`, *optional*, defaults to True): enable kv cache to speed up the inference
-            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.
-        Examples:
-
-        Returns:
-            A list with the generated images.
-        """
-        assert height%16 == 0 and width%16 == 0
-        if self.separate_cfg_infer:
-            self.use_kv_cache = False
-            # raise "Currently, don't support both self.use_kv_cache and self.separate_cfg_infer"
-        if input_images is None:
-            use_img_guidance = False
-        if isinstance(prompt, str):
-            prompt = [prompt]
-            input_images = [input_images] if input_images is not None else None
-
-        input_data = self.processor(prompt, input_images, height=height, width=width, use_img_cfg=use_img_guidance, separate_cfg_input=self.separate_cfg_infer)
-
-        num_prompt = len(prompt)
-        num_cfg = 2 if use_img_guidance else 1
-        latent_size_h, latent_size_w = height//8, width//8
-
-        if seed is not None:
-            generator = torch.Generator(device=self.device).manual_seed(int(seed))
-        else:
-            generator = None
-        latents = torch.randn(num_prompt, 4, latent_size_h, latent_size_w, device=self.device, generator=generator)
-        latents = torch.cat([latents]*(1+num_cfg), 0).to(self.dtype)
-
-        input_img_latents = []
-        if self.separate_cfg_infer:
-            for temp_pixel_values in input_data['input_pixel_values']:
-                temp_input_latents = []
-                for img in temp_pixel_values:
-                    img = self.vae_encode(img.to(self.device), self.dtype)
-                    temp_input_latents.append(img)
-                input_img_latents.append(temp_input_latents)
-        else:
-            for img in input_data['input_pixel_values']:
-                img = self.vae_encode(img.to(self.device), self.dtype)
-                input_img_latents.append(img)
-
-        model_kwargs = dict(input_ids=self.move_to_device(input_data['input_ids']),
-            input_img_latents=input_img_latents,
-            input_image_sizes=input_data['input_image_sizes'],
-            attention_mask=self.move_to_device(input_data["attention_mask"]),
-            position_ids=self.move_to_device(input_data["position_ids"]),
-            cfg_scale=guidance_scale,
-            img_cfg_scale=img_guidance_scale,
-            use_img_cfg=use_img_guidance,
-            use_kv_cache=self.use_kv_cache)
-
-        if self.separate_cfg_infer:
-            func = self.model.forward_with_separate_cfg
-        else:
-            func = self.model.forward_with_cfg
-        self.model.to(self.dtype)
-
-        scheduler = OmniGenScheduler(num_steps=num_inference_steps)
-        samples = scheduler(latents, func, model_kwargs, use_kv_cache=self.use_kv_cache)
-        samples = samples.chunk((1+num_cfg), dim=0)[0]
-
-        if output_type == 'latent':
-            output_images = { 'images': samples }
-            return output_images
-
-        samples = samples.to(self.vae.dtype)
-        if self.vae.config.shift_factor is not None:
-            samples = samples / self.vae.config.scaling_factor + self.vae.config.shift_factor
-        else:
-            samples = samples / self.vae.config.scaling_factor
-        samples = self.vae.decode(samples).sample
-
-        output_samples = (samples * 0.5 + 0.5).clamp(0, 1)*255
-        output_samples = output_samples.permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
-        output_images = []
-        for _i, sample in enumerate(output_samples):
-            output_images.append(Image.fromarray(sample))
-
-        return output_images
diff --git a/modules/omnigen/processor.py b/modules/omnigen/processor.py
deleted file mode 100644
index ada813a8b..000000000
--- a/modules/omnigen/processor.py
+++ /dev/null
@@ -1,312 +0,0 @@
-import os
-import re
-from typing import Dict, List
-import torch
-from torchvision import transforms
-from transformers import AutoTokenizer
-from huggingface_hub import snapshot_download
-from .utils import crop_arr
-
-
-class OmniGenProcessor:
-    def __init__(self,
-                text_tokenizer,
-                max_image_size: int=1024):
-        self.text_tokenizer = text_tokenizer
-        self.max_image_size = max_image_size
-
-        self.image_transform = transforms.Compose([
-            transforms.Lambda(lambda pil_image: crop_arr(pil_image, max_image_size)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
-        ])
-
-        self.collator = OmniGenCollator()
-        self.separate_collator = OmniGenSeparateCollator()
-
-    @classmethod
-    def from_pretrained(cls, model_name):
-        if not os.path.exists(model_name):
-            cache_folder = os.getenv('HF_HUB_CACHE')
-            model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_folder,
-                                           allow_patterns="*.json")
-        text_tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-        return cls(text_tokenizer)
-
-
-    def process_image(self, image):
-        return self.image_transform(image)
-
-    def process_multi_modal_prompt(self, text, input_images):
-        text = self.add_prefix_instruction(text)
-        if input_images is None or len(input_images) == 0:
-            model_inputs = self.text_tokenizer(text)
-            return {"input_ids": model_inputs.input_ids, "pixel_values": None, "image_sizes": None}
-
-        pattern = r"<\|image_\d+\|>"
-        prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)]
-
-        for i in range(1, len(prompt_chunks)):
-            if prompt_chunks[i][0] == 1:
-                prompt_chunks[i] = prompt_chunks[i][1:]
-
-        image_tags = re.findall(pattern, text)
-        image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags]
-
-        unique_image_ids = sorted(list(set(image_ids)))
-        assert unique_image_ids == list(range(1, len(unique_image_ids)+1)), f"image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be {unique_image_ids}"
-        # total images must be the same as the number of image tags
-        assert len(unique_image_ids) == len(input_images), f"total images must be the same as the number of image tags, got {len(unique_image_ids)} image tags and {len(input_images)} images"
-
-        input_images = [input_images[x-1] for x in image_ids]
-
-        all_input_ids = []
-        img_inx = []
-        _idx = 0
-        for i in range(len(prompt_chunks)):
-            all_input_ids.extend(prompt_chunks[i])
-            if i != len(prompt_chunks) -1:
-                start_inx = len(all_input_ids)
-                size = input_images[i].size(-2) *  input_images[i].size(-1) // 16 // 16
-                img_inx.append([start_inx, start_inx+size])
-                all_input_ids.extend([0]*size)
-
-        return {"input_ids": all_input_ids, "pixel_values": input_images, "image_sizes": img_inx}
-
-
-    def add_prefix_instruction(self, prompt):
-        user_prompt = '<|user|>\n'
-        generation_prompt = 'Generate an image according to the following instructions\n'
-        assistant_prompt = '<|assistant|>\n<|diffusion|>'
-        prompt_suffix = "<|end|>\n"
-        prompt = f"{user_prompt}{generation_prompt}{prompt}{prompt_suffix}{assistant_prompt}"
-        return prompt
-
-
-    def __call__(self,
-                instructions: List[str],
-                input_images: List[List[str]] = None,
-                height: int = 1024,
-                width: int = 1024,
-                negative_prompt: str = "low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers.",
-                use_img_cfg: bool = True,
-                separate_cfg_input: bool = False,
-                ) -> Dict:
-
-        if input_images is None:
-            use_img_cfg = False
-        if isinstance(instructions, str):
-            instructions = [instructions]
-            input_images = [input_images]
-
-        input_data = []
-        for i in range(len(instructions)):
-            cur_instruction = instructions[i]
-            cur_input_images = None if input_images is None else input_images[i]
-            if cur_input_images is not None and len(cur_input_images) > 0:
-                cur_input_images = [self.process_image(x) for x in cur_input_images]
-            else:
-                cur_input_images = None
-                assert "<img><|image_1|></img>" not in cur_instruction
-
-            mllm_input = self.process_multi_modal_prompt(cur_instruction, cur_input_images)
-
-            neg_mllm_input, img_cfg_mllm_input = None, None
-            neg_mllm_input = self.process_multi_modal_prompt(negative_prompt, None)
-            if use_img_cfg:
-                if cur_input_images is not None and len(cur_input_images) >= 1:
-                    img_cfg_prompt = [f"<img><|image_{i+1}|></img>" for i in range(len(cur_input_images))]
-                    img_cfg_mllm_input = self.process_multi_modal_prompt(" ".join(img_cfg_prompt), cur_input_images)
-                else:
-                    img_cfg_mllm_input = neg_mllm_input
-
-            input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width]))
-
-        if separate_cfg_input:
-            return self.separate_collator(input_data)
-        return self.collator(input_data)
-
-
-class OmniGenCollator:
-    def __init__(self, pad_token_id=2, hidden_size=3072):
-        self.pad_token_id = pad_token_id
-        self.hidden_size = hidden_size
-
-    def create_position(self, attention_mask, num_tokens_for_output_images):
-        position_ids = []
-        text_length = attention_mask.size(-1)
-        img_length = max(num_tokens_for_output_images)
-        for mask in attention_mask:
-            temp_l = torch.sum(mask)
-            temp_position = [0]*(text_length-temp_l) + [i for i in range(temp_l+img_length+1)] # we add a time embedding into the sequence, so add one more token
-            position_ids.append(temp_position)
-        return torch.LongTensor(position_ids)
-
-    def create_mask(self, attention_mask, num_tokens_for_output_images):
-        extended_mask = []
-        padding_images = []
-        text_length = attention_mask.size(-1)
-        img_length = max(num_tokens_for_output_images)
-        seq_len = text_length + img_length + 1 # we add a time embedding into the sequence, so add one more token
-        inx = 0
-        for mask in attention_mask:
-            temp_l = torch.sum(mask)
-            pad_l = text_length - temp_l
-
-            temp_mask = torch.tril(torch.ones(size=(temp_l+1, temp_l+1)))
-
-            image_mask = torch.zeros(size=(temp_l+1, img_length))
-            temp_mask = torch.cat([temp_mask, image_mask], dim=-1)
-
-            image_mask = torch.ones(size=(img_length, temp_l+img_length+1))
-            temp_mask = torch.cat([temp_mask, image_mask], dim=0)
-
-            if pad_l > 0:
-                pad_mask = torch.zeros(size=(temp_l+1+img_length, pad_l))
-                temp_mask = torch.cat([pad_mask, temp_mask], dim=-1)
-
-                pad_mask = torch.ones(size=(pad_l, seq_len))
-                temp_mask = torch.cat([pad_mask, temp_mask], dim=0)
-
-            true_img_length = num_tokens_for_output_images[inx]
-            pad_img_length = img_length - true_img_length
-            if pad_img_length > 0:
-                temp_mask[:, -pad_img_length:] = 0
-                temp_padding_imgs = torch.zeros(size=(1, pad_img_length, self.hidden_size))
-            else:
-                temp_padding_imgs = None
-
-            extended_mask.append(temp_mask.unsqueeze(0))
-            padding_images.append(temp_padding_imgs)
-            inx += 1
-        return torch.cat(extended_mask, dim=0), padding_images
-
-    def adjust_attention_for_input_images(self, attention_mask, image_sizes):
-        for b_inx in image_sizes.keys():
-            for start_inx, end_inx in image_sizes[b_inx]:
-                attention_mask[b_inx][start_inx:end_inx, start_inx:end_inx] = 1
-
-        return attention_mask
-
-    def pad_input_ids(self, input_ids, image_sizes):
-        max_l = max([len(x) for x in input_ids])
-        padded_ids = []
-        attention_mask = []
-        _new_image_sizes = []
-
-        for i in range(len(input_ids)):
-            temp_ids = input_ids[i]
-            temp_l = len(temp_ids)
-            pad_l = max_l - temp_l
-            if pad_l == 0:
-                attention_mask.append([1]*max_l)
-                padded_ids.append(temp_ids)
-            else:
-                attention_mask.append([0]*pad_l+[1]*temp_l)
-                padded_ids.append([self.pad_token_id]*pad_l+temp_ids)
-
-            if i in image_sizes:
-                new_inx = []
-                for old_inx in image_sizes[i]:
-                    new_inx.append([x+pad_l for x in old_inx])
-                image_sizes[i] = new_inx
-
-        return torch.LongTensor(padded_ids), torch.LongTensor(attention_mask), image_sizes
-
-
-    def process_mllm_input(self, mllm_inputs, target_img_size):
-        num_tokens_for_output_images = []
-        for img_size in target_img_size:
-            num_tokens_for_output_images.append(img_size[0]*img_size[1]//16//16)
-
-        pixel_values, image_sizes = [], {}
-        b_inx = 0
-        for x in mllm_inputs:
-            if x['pixel_values'] is not None:
-                pixel_values.extend(x['pixel_values'])
-                for size in x['image_sizes']:
-                    if b_inx not in image_sizes:
-                        image_sizes[b_inx] = [size]
-                    else:
-                        image_sizes[b_inx].append(size)
-            b_inx += 1
-        pixel_values = [x.unsqueeze(0) for x in pixel_values]
-
-        input_ids = [x['input_ids'] for x in mllm_inputs]
-        padded_input_ids, attention_mask, image_sizes = self.pad_input_ids(input_ids, image_sizes)
-        position_ids = self.create_position(attention_mask, num_tokens_for_output_images)
-        attention_mask, padding_images = self.create_mask(attention_mask, num_tokens_for_output_images)
-        attention_mask = self.adjust_attention_for_input_images(attention_mask, image_sizes)
-
-        return padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes
-
-    def __call__(self, features):
-        mllm_inputs = [f[0] for f in features]
-        cfg_mllm_inputs = [f[1] for f in features]
-        img_cfg_mllm_input = [f[2] for f in features]
-        target_img_size = [f[3] for f in features]
-
-        if img_cfg_mllm_input[0] is not None:
-            mllm_inputs = mllm_inputs + cfg_mllm_inputs + img_cfg_mllm_input
-            target_img_size = target_img_size + target_img_size + target_img_size
-        else:
-            mllm_inputs = mllm_inputs + cfg_mllm_inputs
-            target_img_size = target_img_size + target_img_size
-
-
-        all_padded_input_ids, all_position_ids, all_attention_mask, all_padding_images, all_pixel_values, all_image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
-
-        data = {"input_ids": all_padded_input_ids,
-        "attention_mask": all_attention_mask,
-        "position_ids": all_position_ids,
-        "input_pixel_values": all_pixel_values,
-        "input_image_sizes": all_image_sizes,
-        "padding_images": all_padding_images,
-        }
-        return data
-
-
-class OmniGenSeparateCollator(OmniGenCollator):
-    def __call__(self, features):
-        mllm_inputs = [f[0] for f in features]
-        cfg_mllm_inputs = [f[1] for f in features]
-        img_cfg_mllm_input = [f[2] for f in features]
-        target_img_size = [f[3] for f in features]
-
-        all_padded_input_ids, all_attention_mask, all_position_ids, all_pixel_values, all_image_sizes, all_padding_images = [], [], [], [], [], []
-
-        padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
-        all_padded_input_ids.append(padded_input_ids)
-        all_attention_mask.append(attention_mask)
-        all_position_ids.append(position_ids)
-        all_pixel_values.append(pixel_values)
-        all_image_sizes.append(image_sizes)
-        all_padding_images.append(padding_images)
-
-        if cfg_mllm_inputs[0] is not None:
-            padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(cfg_mllm_inputs, target_img_size)
-            all_padded_input_ids.append(padded_input_ids)
-            all_attention_mask.append(attention_mask)
-            all_position_ids.append(position_ids)
-            all_pixel_values.append(pixel_values)
-            all_image_sizes.append(image_sizes)
-            all_padding_images.append(padding_images)
-        if img_cfg_mllm_input[0] is not None:
-            padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(img_cfg_mllm_input, target_img_size)
-            all_padded_input_ids.append(padded_input_ids)
-            all_attention_mask.append(attention_mask)
-            all_position_ids.append(position_ids)
-            all_pixel_values.append(pixel_values)
-            all_image_sizes.append(image_sizes)
-            all_padding_images.append(padding_images)
-
-        data = {"input_ids": all_padded_input_ids,
-        "attention_mask": all_attention_mask,
-        "position_ids": all_position_ids,
-        "input_pixel_values": all_pixel_values,
-        "input_image_sizes": all_image_sizes,
-        "padding_images": all_padding_images,
-        }
-        return data
diff --git a/modules/omnigen/scheduler.py b/modules/omnigen/scheduler.py
deleted file mode 100644
index 0764fd8f0..000000000
--- a/modules/omnigen/scheduler.py
+++ /dev/null
@@ -1,55 +0,0 @@
-import torch
-from tqdm import tqdm
-from transformers.cache_utils import Cache, DynamicCache
-
-class OmniGenScheduler:
-    def __init__(self, num_steps: int=50, time_shifting_factor: int=1):
-        self.num_steps = num_steps
-        self.time_shift = time_shifting_factor
-
-        t = torch.linspace(0, 1, num_steps+1)
-        t = t / (t + time_shifting_factor - time_shifting_factor * t)
-        self.sigma = t
-    
-    def crop_kv_cache(self, past_key_values, num_tokens_for_img):
-        crop_past_key_values = ()
-        for layer_idx in range(len(past_key_values)):
-            key_states, value_states = past_key_values[layer_idx][:2]
-            crop_past_key_values += ((key_states[..., :-(num_tokens_for_img+1), :], value_states[..., :-(num_tokens_for_img+1), :], ),)
-        return crop_past_key_values
-        # return DynamicCache.from_legacy_cache(crop_past_key_values)
-
-    def crop_position_ids_for_cache(self, position_ids, num_tokens_for_img):
-        if isinstance(position_ids, list):
-            for i in range(len(position_ids)):
-                position_ids[i] = position_ids[i][:, -(num_tokens_for_img+1):]
-        else:
-            position_ids = position_ids[:, -(num_tokens_for_img+1):]
-        return position_ids
-
-    def crop_attention_mask_for_cache(self, attention_mask, num_tokens_for_img):
-        if isinstance(attention_mask, list):
-            return [x[..., -(num_tokens_for_img+1):, :] for x in attention_mask]
-        return attention_mask[..., -(num_tokens_for_img+1):, :]
-
-    def __call__(self, z, func, model_kwargs, use_kv_cache: bool=True):
-        past_key_values = None
-        for i in tqdm(range(self.num_steps)):
-            timesteps = torch.zeros(size=(len(z), )).to(z.device) + self.sigma[i]
-            pred, temp_past_key_values = func(z, timesteps, past_key_values=past_key_values, **model_kwargs)
-            sigma_next = self.sigma[i+1]
-            sigma = self.sigma[i]
-            z = z + (sigma_next - sigma) * pred
-            if i == 0 and use_kv_cache:
-                num_tokens_for_img = z.size(-1)*z.size(-2) // 4
-                if isinstance(temp_past_key_values, list):
-                    past_key_values = [self.crop_kv_cache(x, num_tokens_for_img) for x in temp_past_key_values]
-                    model_kwargs['input_ids'] = [None] * len(temp_past_key_values)
-                else:
-                    past_key_values = self.crop_kv_cache(temp_past_key_values, num_tokens_for_img)
-                    model_kwargs['input_ids'] = None
-
-                model_kwargs['position_ids'] = self.crop_position_ids_for_cache(model_kwargs['position_ids'], num_tokens_for_img)
-                model_kwargs['attention_mask'] = self.crop_attention_mask_for_cache(model_kwargs['attention_mask'], num_tokens_for_img)
-        return z
-
diff --git a/modules/omnigen/transformer.py b/modules/omnigen/transformer.py
deleted file mode 100644
index d166309ca..000000000
--- a/modules/omnigen/transformer.py
+++ /dev/null
@@ -1,164 +0,0 @@
-import math
-import warnings
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-from huggingface_hub import snapshot_download
-
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-    SequenceClassifierOutputWithPast,
-    TokenClassifierOutput,
-)
-from transformers.modeling_utils import PreTrainedModel
-from transformers import Phi3Config, Phi3Model
-from transformers.cache_utils import Cache, DynamicCache, StaticCache
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-class Phi3Transformer(Phi3Model):
-    """
-    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]
-    We only modified the attention mask
-    Args:
-        config: Phi3Config
-    """
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-    ) -> Union[Tuple, BaseModelOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning_once(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
-        # kept for BC (non `Cache` `past_key_values` inputs)
-        return_legacy_cache = False
-        if use_cache and not isinstance(past_key_values, Cache):
-            return_legacy_cache = True
-            if past_key_values is None:
-                past_key_values = DynamicCache()
-            else:
-                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-                logger.warning_once(
-                    "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
-                    "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
-                    "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
-                )
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        if attention_mask is not None and attention_mask.dim() == 3:
-            dtype = inputs_embeds.dtype
-            min_dtype = torch.finfo(dtype).min
-            attention_mask = (1 - attention_mask) * min_dtype
-            attention_mask = attention_mask.unsqueeze(1).to(inputs_embeds.dtype)
-        else:
-            raise 
-            # causal_mask = self._update_causal_mask(
-            #     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
-            # )
-
-        hidden_states = inputs_embeds
-
-        # create position embeddings to be shared across the decoder layers
-        position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-        next_decoder_cache = None
-
-        for decoder_layer in self.layers:
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            if self.gradient_checkpointing and self.training:
-                layer_outputs = self._gradient_checkpointing_func(
-                    decoder_layer.__call__,
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    past_key_values,
-                    output_attentions,
-                    use_cache,
-                    cache_position,
-                    position_embeddings,
-                )
-            else:
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_values,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                    cache_position=cache_position,
-                    position_embeddings=position_embeddings,
-                )
-
-            hidden_states = layer_outputs[0]
-
-            if use_cache:
-                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        next_cache = next_decoder_cache if use_cache else None
-        if return_legacy_cache:
-            next_cache = next_cache.to_legacy_cache()
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-
diff --git a/modules/omnigen/utils.py b/modules/omnigen/utils.py
deleted file mode 100644
index bf0a6de62..000000000
--- a/modules/omnigen/utils.py
+++ /dev/null
@@ -1,105 +0,0 @@
-import logging
-
-from PIL import Image
-import torch
-import numpy as np
-
-def create_logger(logging_dir):
-    """
-    Create a logger that writes to a log file and stdout.
-    """
-    logging.basicConfig(
-        level=logging.INFO,
-        format='[\033[34m%(asctime)s\033[0m] %(message)s',
-        datefmt='%Y-%m-%d %H:%M:%S',
-        handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
-    )
-    logger = logging.getLogger(__name__)
-    return logger
-
-
-@torch.no_grad()
-def update_ema(ema_model, model, decay=0.9999):
-    """
-    Step the EMA model towards the current model.
-    """
-    ema_params = dict(ema_model.named_parameters())
-    for name, param in model.named_parameters():
-        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
-
-
-def requires_grad(model, flag=True):
-    """
-    Set requires_grad flag for all parameters in a model.
-    """
-    for p in model.parameters():
-        p.requires_grad = flag
-
-
-def center_crop_arr(pil_image, image_size):
-    """
-    Center cropping implementation from ADM.
-    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
-    """
-    while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(
-            tuple(x // 2 for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(
-        tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS
-    )
-
-    arr = np.array(pil_image)
-    crop_y = (arr.shape[0] - image_size) // 2
-    crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
-
-
-def crop_arr(pil_image, max_image_size):
-    while min(*pil_image.size) >= 2 * max_image_size:
-        pil_image = pil_image.resize(
-            tuple(x // 2 for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    if max(*pil_image.size) > max_image_size:
-        scale = max_image_size / max(*pil_image.size)
-        pil_image = pil_image.resize(
-            tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    if min(*pil_image.size) < 16:
-        scale = 16 / min(*pil_image.size)
-        pil_image = pil_image.resize(
-            tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    arr = np.array(pil_image)
-    crop_y1 = (arr.shape[0] % 16) // 2
-    crop_y2 = arr.shape[0] % 16 - crop_y1
-
-    crop_x1 = (arr.shape[1] % 16) // 2
-    crop_x2 = arr.shape[1] % 16 - crop_x1
-
-    arr = arr[crop_y1:arr.shape[0]-crop_y2, crop_x1:arr.shape[1]-crop_x2]
-    return Image.fromarray(arr)
-
-
-def vae_encode(vae, x, weight_dtype):
-    if x is not None:
-        if vae.config.shift_factor is not None:
-            x = vae.encode(x).latent_dist.sample()
-            x = (x - vae.config.shift_factor) * vae.config.scaling_factor
-        else:
-            x = vae.encode(x).latent_dist.sample().mul_(vae.config.scaling_factor)
-        x = x.to(weight_dtype)
-    return x
-
-
-def vae_encode_list(vae, x, weight_dtype):
-    latents = []
-    for img in x:
-        img = vae_encode(vae, img, weight_dtype)
-        latents.append(img)
-    return latents
diff --git a/modules/omnigen2/__init__.py b/modules/omnigen2/__init__.py
new file mode 100644
index 000000000..82d539e60
--- /dev/null
+++ b/modules/omnigen2/__init__.py
@@ -0,0 +1,3 @@
+from transformers import Qwen2_5_VLForConditionalGeneration
+from .pipeline_omnigen2 import OmniGen2Pipeline
+from .models.transformers import OmniGen2Transformer2DModel
diff --git a/modules/omnigen2/image_processor.py b/modules/omnigen2/image_processor.py
new file mode 100644
index 000000000..00b1fb3ed
--- /dev/null
+++ b/modules/omnigen2/image_processor.py
@@ -0,0 +1,265 @@
+# Copyright 2024 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.
+
+import warnings
+from typing import Optional, Tuple, Union
+
+import numpy as np
+import PIL.Image
+import torch
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor, is_valid_image_imagelist
+from diffusers.configuration_utils import register_to_config
+
+class OmniGen2ImageProcessor(VaeImageProcessor):
+    """
+    Image processor for PixArt image resize and crop.
+
+    Args:
+        do_resize (`bool`, *optional*, defaults to `True`):
+            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
+            `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
+        vae_scale_factor (`int`, *optional*, defaults to `8`):
+            VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
+        resample (`str`, *optional*, defaults to `lanczos`):
+            Resampling filter to use when resizing the image.
+        do_normalize (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the image to [-1,1].
+        do_binarize (`bool`, *optional*, defaults to `False`):
+            Whether to binarize the image to 0/1.
+        do_convert_rgb (`bool`, *optional*, defaults to be `False`):
+            Whether to convert the images to RGB format.
+        do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
+            Whether to convert the images to grayscale format.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        do_resize: bool = True,
+        vae_scale_factor: int = 16,
+        resample: str = "lanczos",
+        max_pixels: Optional[int] = None,
+        max_side_length: Optional[int] = None,
+        do_normalize: bool = True,
+        do_binarize: bool = False,
+        do_convert_grayscale: bool = False,
+    ):
+        super().__init__(
+            do_resize=do_resize,
+            vae_scale_factor=vae_scale_factor,
+            resample=resample,
+            do_normalize=do_normalize,
+            do_binarize=do_binarize,
+            do_convert_grayscale=do_convert_grayscale,
+        )
+
+        self.max_pixels = max_pixels
+        self.max_side_length = max_side_length
+
+    def get_new_height_width(
+        self,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor],
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        max_pixels: Optional[int] = None,
+        max_side_length: Optional[int] = None,
+    ) -> Tuple[int, int]:
+        r"""
+        Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
+
+        Args:
+            image (`Union[PIL.Image.Image, np.ndarray, torch.Tensor]`):
+                The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it
+                should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch
+                tensor, it should have shape `[batch, channels, height, width]`.
+            height (`Optional[int]`, *optional*, defaults to `None`):
+                The height of the preprocessed image. If `None`, the height of the `image` input will be used.
+            width (`Optional[int]`, *optional*, defaults to `None`):
+                The width of the preprocessed image. If `None`, the width of the `image` input will be used.
+
+        Returns:
+            `Tuple[int, int]`:
+                A tuple containing the height and width, both resized to the nearest integer multiple of
+                `vae_scale_factor`.
+        """
+
+        if height is None:
+            if isinstance(image, PIL.Image.Image):
+                height = image.height
+            elif isinstance(image, torch.Tensor):
+                height = image.shape[2]
+            else:
+                height = image.shape[1]
+
+        if width is None:
+            if isinstance(image, PIL.Image.Image):
+                width = image.width
+            elif isinstance(image, torch.Tensor):
+                width = image.shape[3]
+            else:
+                width = image.shape[2]
+
+        if max_side_length is None:
+            max_side_length = self.max_side_length
+
+        if max_pixels is None:
+            max_pixels = self.max_pixels
+
+        ratio = 1.0
+        if max_side_length is not None:
+            if height > width:
+                max_side_length_ratio = max_side_length / height
+            else:
+                max_side_length_ratio = max_side_length / width
+
+        cur_pixels = height * width
+        max_pixels_ratio = (max_pixels / cur_pixels) ** 0.5
+        ratio = min(max_pixels_ratio, max_side_length_ratio, 1.0) # do not upscale input image
+
+        new_height, new_width = int(height * ratio) // self.config.vae_scale_factor * self.config.vae_scale_factor, int(width * ratio) // self.config.vae_scale_factor * self.config.vae_scale_factor
+        return new_height, new_width
+
+    def preprocess(
+        self,
+        image: PipelineImageInput,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        max_pixels: Optional[int] = None,
+        max_side_length: Optional[int] = None,
+        resize_mode: str = "default",  # "default", "fill", "crop"
+        crops_coords: Optional[Tuple[int, int, int, int]] = None,
+    ) -> torch.Tensor:
+        """
+        Preprocess the image input.
+
+        Args:
+            image (`PipelineImageInput`):
+                The image input, accepted formats are PIL images, NumPy arrays, PyTorch tensors; Also accept list of
+                supported formats.
+            height (`int`, *optional*):
+                The height in preprocessed image. If `None`, will use the `get_default_height_width()` to get default
+                height.
+            width (`int`, *optional*):
+                The width in preprocessed. If `None`, will use get_default_height_width()` to get the default width.
+            resize_mode (`str`, *optional*, defaults to `default`):
+                The resize mode, can be one of `default` or `fill`. If `default`, will resize the image to fit within
+                the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, will
+                resize the image to fit within the specified width and height, maintaining the aspect ratio, and then
+                center the image within the dimensions, filling empty with data from image. If `crop`, will resize the
+                image to fit within the specified width and height, maintaining the aspect ratio, and then center the
+                image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
+                supported for PIL image input.
+            crops_coords (`List[Tuple[int, int, int, int]]`, *optional*, defaults to `None`):
+                The crop coordinates for each image in the batch. If `None`, will not crop the image.
+
+        Returns:
+            `torch.Tensor`:
+                The preprocessed image.
+        """
+        supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
+
+        # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
+        if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3:
+            if isinstance(image, torch.Tensor):
+                # if image is a pytorch tensor could have 2 possible shapes:
+                #    1. batch x height x width: we should insert the channel dimension at position 1
+                #    2. channel x height x width: we should insert batch dimension at position 0,
+                #       however, since both channel and batch dimension has same size 1, it is same to insert at position 1
+                #    for simplicity, we insert a dimension of size 1 at position 1 for both cases
+                image = image.unsqueeze(1)
+            else:
+                # if it is a numpy array, it could have 2 possible shapes:
+                #   1. batch x height x width: insert channel dimension on last position
+                #   2. height x width x channel: insert batch dimension on first position
+                if image.shape[-1] == 1:
+                    image = np.expand_dims(image, axis=0)
+                else:
+                    image = np.expand_dims(image, axis=-1)
+
+        if isinstance(image, list) and isinstance(image[0], np.ndarray) and image[0].ndim == 4:
+            warnings.warn(
+                "Passing `image` as a list of 4d np.ndarray is deprecated."
+                "Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray",
+                FutureWarning,
+            )
+            image = np.concatenate(image, axis=0)
+        if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4:
+            warnings.warn(
+                "Passing `image` as a list of 4d torch.Tensor is deprecated."
+                "Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor",
+                FutureWarning,
+            )
+            image = torch.cat(image, axis=0)
+
+        if not is_valid_image_imagelist(image):
+            raise ValueError(
+                f"Input is in incorrect format. Currently, we only support {', '.join(str(x) for x in supported_formats)}"
+            )
+        if not isinstance(image, list):
+            image = [image]
+
+        if isinstance(image[0], PIL.Image.Image):
+            if crops_coords is not None:
+                image = [i.crop(crops_coords) for i in image]
+            if self.config.do_resize:
+                height, width = self.get_new_height_width(image[0], height, width, max_pixels, max_side_length)
+                image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image]
+            if self.config.do_convert_rgb:
+                image = [self.convert_to_rgb(i) for i in image]
+            elif self.config.do_convert_grayscale:
+                image = [self.convert_to_grayscale(i) for i in image]
+            image = self.pil_to_numpy(image)  # to np
+            image = self.numpy_to_pt(image)  # to pt
+
+        elif isinstance(image[0], np.ndarray):
+            image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)
+
+            image = self.numpy_to_pt(image)
+
+            height, width = self.get_new_height_width(image, height, width, max_pixels, max_side_length)
+            if self.config.do_resize:
+                image = self.resize(image, height, width)
+
+        elif isinstance(image[0], torch.Tensor):
+            image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
+
+            if self.config.do_convert_grayscale and image.ndim == 3:
+                image = image.unsqueeze(1)
+
+            channel = image.shape[1]
+            # don't need any preprocess if the image is latents
+            if channel == self.config.vae_latent_channels:
+                return image
+
+            height, width = self.get_new_height_width(image, height, width, max_pixels, max_side_length)
+            if self.config.do_resize:
+                image = self.resize(image, height, width)
+
+        # expected range [0,1], normalize to [-1,1]
+        do_normalize = self.config.do_normalize
+        if do_normalize and image.min() < 0:
+            warnings.warn(
+                "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
+                f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",
+                FutureWarning,
+            )
+            do_normalize = False
+        if do_normalize:
+            image = self.normalize(image)
+
+        if self.config.do_binarize:
+            image = self.binarize(image)
+
+        return image
diff --git a/modules/omnigen2/models/attention_processor.py b/modules/omnigen2/models/attention_processor.py
new file mode 100644
index 000000000..ffeb8e141
--- /dev/null
+++ b/modules/omnigen2/models/attention_processor.py
@@ -0,0 +1,141 @@
+"""
+OmniGen2 Attention Processor Module
+
+Copyright 2025 BAAI, The OmniGen2 Team and 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.
+"""
+
+import warnings
+import math
+from typing import Optional, Tuple, Dict, Any
+
+import torch
+import torch.nn.functional as F
+from einops import repeat
+
+from diffusers.models.attention_processor import Attention
+from .embeddings import apply_rotary_emb
+
+
+class OmniGen2AttnProcessor:
+    """
+    Processor for implementing scaled dot-product attention with flash attention and variable length sequences.
+
+    This processor is optimized for PyTorch 2.0 and implements:
+    - Flash attention with variable length sequences
+    - Rotary position embeddings (RoPE)
+    - Query-Key normalization
+    - Proportional attention scaling
+
+    Args:
+        None
+
+    Raises:
+        ImportError: If PyTorch version is less than 2.0
+    """
+
+    def __init__(self) -> None:
+        """Initialize the attention processor."""
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "OmniGen2AttnProcessor requires PyTorch 2.0. "
+                "Please upgrade PyTorch to version 2.0 or later."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        base_sequence_length: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Process attention computation with flash attention.
+
+        Args:
+            attn: Attention module
+            hidden_states: Hidden states tensor of shape (batch_size, seq_len, hidden_dim)
+            encoder_hidden_states: Encoder hidden states tensor
+            attention_mask: Optional attention mask tensor
+            image_rotary_emb: Optional rotary embeddings for image tokens
+            base_sequence_length: Optional base sequence length for proportional attention
+
+        Returns:
+            torch.Tensor: Processed hidden states after attention computation
+        """
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        # Get Query-Key-Value Pair
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query_dim = query.shape[-1]
+        inner_dim = key.shape[-1]
+        head_dim = query_dim // attn.heads
+        dtype = query.dtype
+
+        # Get key-value heads
+        kv_heads = inner_dim // head_dim
+
+        # Reshape tensors for attention computation
+        query = query.view(batch_size, -1, attn.heads, head_dim)
+        key = key.view(batch_size, -1, kv_heads, head_dim)
+        value = value.view(batch_size, -1, kv_heads, head_dim)
+
+        # Apply Query-Key normalization
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply Rotary Position Embeddings
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, use_real=False)
+            key = apply_rotary_emb(key, image_rotary_emb, use_real=False)
+
+        query, key = query.to(dtype), key.to(dtype)
+
+        # Calculate attention scale
+        if base_sequence_length is not None:
+            softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
+        else:
+            softmax_scale = attn.scale
+
+        # scaled_dot_product_attention expects attention_mask shape to be
+        # (batch, heads, source_length, target_length)
+        if attention_mask is not None:
+            attention_mask = attention_mask.bool().view(batch_size, 1, 1, -1)
+
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        # explicitly repeat key and value to match query length, otherwise using enable_gqa=True results in MATH backend of sdpa in our test of pytorch2.6
+        key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
+        value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, scale=softmax_scale
+        )
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.type_as(query)
+
+        # Apply output projection
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
diff --git a/modules/omnigen2/models/embeddings.py b/modules/omnigen2/models/embeddings.py
new file mode 100644
index 000000000..574ff01bd
--- /dev/null
+++ b/modules/omnigen2/models/embeddings.py
@@ -0,0 +1,99 @@
+# Copyright 2024 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 typing import List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+from modules import devices
+
+
+# Omnigen uses x.shape[-1] // 2 instead of -1
+# Functionally the same but -1 does fail with when the shape becomes 0
+if devices.backend != "ipex":
+    def apply_rotary_emb(
+        x: torch.Tensor,
+        freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+        use_real: bool = True,
+        use_real_unbind_dim: int = -1,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+        to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+        reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+        tensors contain rotary embeddings and are returned as real tensors.
+
+        Args:
+            x (`torch.Tensor`):
+                Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+            freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+        """
+        if use_real:
+            cos, sin = freqs_cis  # [S, D]
+            cos = cos[None, None]
+            sin = sin[None, None]
+            cos, sin = cos.to(x.device), sin.to(x.device)
+
+            if use_real_unbind_dim == -1:
+                # Used for flux, cogvideox, hunyuan-dit
+                x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+                x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+            elif use_real_unbind_dim == -2:
+                # Used for Stable Audio, OmniGen and CogView4
+                x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+                x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+            else:
+                raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+            out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+            return out
+        else:
+            # used for lumina
+            # x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+            x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], x.shape[-1] // 2, 2))
+            freqs_cis = freqs_cis.unsqueeze(2)
+            x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+            return x_out.type_as(x)
+else:
+    def apply_rotary_emb(x, freqs_cis, use_real: bool = True, use_real_unbind_dim: int = -1):
+        if use_real:
+            cos, sin = freqs_cis  # [S, D]
+            cos = cos[None, None]
+            sin = sin[None, None]
+            cos, sin = cos.to(x.device), sin.to(x.device)
+
+            if use_real_unbind_dim == -1:
+                # Used for flux, cogvideox, hunyuan-dit
+                x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+                x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+            elif use_real_unbind_dim == -2:
+                # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+                x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+                x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+            else:
+                raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+            out = (x.to(dtype=torch.float32) * cos + x_rotated.to(dtype=torch.float32) * sin).to(x.dtype)
+            return out
+        else:
+            # used for lumina
+            # force cpu with Alchemist
+            x_rotated = torch.view_as_complex(x.to("cpu").to(dtype=torch.float32).reshape(*x.shape[:-1], x.shape[-1] // 2, 2))
+            freqs_cis = freqs_cis.to("cpu").unsqueeze(2)
+            x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+            return x_out.type_as(x).to(x.device)
diff --git a/modules/omnigen2/models/transformers/__init__.py b/modules/omnigen2/models/transformers/__init__.py
new file mode 100644
index 000000000..b2a23df90
--- /dev/null
+++ b/modules/omnigen2/models/transformers/__init__.py
@@ -0,0 +1,3 @@
+from .transformer_omnigen2 import OmniGen2Transformer2DModel
+
+__all__ = ["OmniGen2Transformer2DModel"]
diff --git a/modules/omnigen2/models/transformers/block_lumina2.py b/modules/omnigen2/models/transformers/block_lumina2.py
new file mode 100644
index 000000000..601736841
--- /dev/null
+++ b/modules/omnigen2/models/transformers/block_lumina2.py
@@ -0,0 +1,64 @@
+
+# Copyright 2024 Alpha-VLLM Authors and 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 typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch.nn import RMSNorm
+from diffusers.models.embeddings import Timesteps, TimestepEmbedding
+
+
+# Makes timestep_scale configurable
+# Omnigen 2 uses timestep_scale=1000
+class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int = 4096,
+        text_feat_dim: int = 2048,
+        frequency_embedding_size: int = 256,
+        norm_eps: float = 1e-5,
+        timestep_scale: float = 1.0,
+    ) -> None:
+        super().__init__()
+
+        self.time_proj = Timesteps(
+            num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=timestep_scale
+        )
+
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024)
+        )
+
+        self.caption_embedder = nn.Sequential(
+            RMSNorm(text_feat_dim, eps=norm_eps),
+            nn.Linear(text_feat_dim, hidden_size, bias=True),
+        )
+
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        nn.init.trunc_normal_(self.caption_embedder[1].weight, std=0.02)
+        nn.init.zeros_(self.caption_embedder[1].bias)
+
+    def forward(
+        self, timestep: torch.Tensor, text_hidden_states: torch.Tensor, dtype: torch.dtype
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        timestep_proj = self.time_proj(timestep).to(dtype=dtype)
+        time_embed = self.timestep_embedder(timestep_proj)
+        caption_embed = self.caption_embedder(text_hidden_states)
+        return time_embed, caption_embed
diff --git a/modules/omnigen2/models/transformers/repo.py b/modules/omnigen2/models/transformers/repo.py
new file mode 100644
index 000000000..8f7c47566
--- /dev/null
+++ b/modules/omnigen2/models/transformers/repo.py
@@ -0,0 +1,129 @@
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+
+from einops import repeat
+from diffusers.models.embeddings import get_1d_rotary_pos_embed
+
+class OmniGen2RotaryPosEmbed(nn.Module):
+    def __init__(self, theta: int,
+                 axes_dim: Tuple[int, int, int],
+                 axes_lens: Tuple[int, int, int] = (300, 512, 512),
+                 patch_size: int = 2):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        self.axes_lens = axes_lens
+        self.patch_size = patch_size
+
+    @staticmethod
+    def get_freqs_cis(axes_dim: Tuple[int, int, int],
+                      axes_lens: Tuple[int, int, int],
+                      theta: int) -> List[torch.Tensor]:
+        freqs_cis = []
+        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
+        for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
+            emb = get_1d_rotary_pos_embed(d, e, theta=theta, freqs_dtype=freqs_dtype)
+            freqs_cis.append(emb)
+        return freqs_cis
+
+    def _get_freqs_cis(self, freqs_cis, ids: torch.Tensor) -> torch.Tensor:
+        device = ids.device
+        if ids.device.type == "mps":
+            ids = ids.to("cpu")
+
+        result = []
+        for i in range(len(self.axes_dim)):
+            freqs = freqs_cis[i].to(ids.device)
+            index = ids[:, :, i : i + 1].repeat(1, 1, freqs.shape[-1]).to(torch.int64)
+            result.append(torch.gather(freqs.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
+        return torch.cat(result, dim=-1).to(device)
+
+    def forward(
+        self,
+        freqs_cis,
+        attention_mask,
+        l_effective_ref_img_len,
+        l_effective_img_len,
+        ref_img_sizes,
+        img_sizes,
+        device
+    ):
+        batch_size = len(attention_mask)
+        p = self.patch_size
+
+        encoder_seq_len = attention_mask.shape[1]
+        l_effective_cap_len = attention_mask.sum(dim=1).tolist()
+
+        seq_lengths = [cap_len + sum(ref_img_len) + img_len for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len)]
+
+        max_seq_len = max(seq_lengths)
+        max_ref_img_len = max([sum(ref_img_len) for ref_img_len in l_effective_ref_img_len])
+        max_img_len = max(l_effective_img_len)
+
+        # Create position IDs
+        position_ids = torch.zeros(batch_size, max_seq_len, 3, dtype=torch.int32, device=device)
+
+        for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)):
+            # add text position ids
+            position_ids[i, :cap_seq_len] = repeat(torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3")
+
+            pe_shift = cap_seq_len
+            pe_shift_len = cap_seq_len
+
+            if ref_img_sizes[i] is not None:
+                for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]):
+                    H, W = ref_img_size
+                    ref_H_tokens, ref_W_tokens = H // p, W // p
+                    assert ref_H_tokens * ref_W_tokens == ref_img_len
+                    # add image position ids
+
+                    row_ids = repeat(torch.arange(ref_H_tokens, dtype=torch.int32, device=device), "h -> h w", w=ref_W_tokens).flatten()
+                    col_ids = repeat(torch.arange(ref_W_tokens, dtype=torch.int32, device=device), "w -> h w", h=ref_H_tokens).flatten()
+                    position_ids[i, pe_shift_len:pe_shift_len + ref_img_len, 0] = pe_shift
+                    position_ids[i, pe_shift_len:pe_shift_len + ref_img_len, 1] = row_ids
+                    position_ids[i, pe_shift_len:pe_shift_len + ref_img_len, 2] = col_ids
+
+                    pe_shift += max(ref_H_tokens, ref_W_tokens)
+                    pe_shift_len += ref_img_len
+
+            H, W = img_sizes[i]
+            H_tokens, W_tokens = H // p, W // p
+            assert H_tokens * W_tokens == l_effective_img_len[i]
+
+            row_ids = repeat(torch.arange(H_tokens, dtype=torch.int32, device=device), "h -> h w", w=W_tokens).flatten()
+            col_ids = repeat(torch.arange(W_tokens, dtype=torch.int32, device=device), "w -> h w", h=H_tokens).flatten()
+
+            assert pe_shift_len + l_effective_img_len[i] == seq_len
+            position_ids[i, pe_shift_len: seq_len, 0] = pe_shift
+            position_ids[i, pe_shift_len: seq_len, 1] = row_ids
+            position_ids[i, pe_shift_len: seq_len, 2] = col_ids
+
+        # Get combined rotary embeddings
+        freqs_cis = self._get_freqs_cis(freqs_cis, position_ids)
+
+        # create separate rotary embeddings for captions and images
+        cap_freqs_cis = torch.zeros(
+            batch_size, encoder_seq_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype
+        )
+        ref_img_freqs_cis = torch.zeros(
+            batch_size, max_ref_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype
+        )
+        img_freqs_cis = torch.zeros(
+            batch_size, max_img_len, freqs_cis.shape[-1], device=device, dtype=freqs_cis.dtype
+        )
+
+        for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths)):
+            cap_freqs_cis[i, :cap_seq_len] = freqs_cis[i, :cap_seq_len]
+            ref_img_freqs_cis[i, :sum(ref_img_len)] = freqs_cis[i, cap_seq_len:cap_seq_len + sum(ref_img_len)]
+            img_freqs_cis[i, :img_len] = freqs_cis[i, cap_seq_len + sum(ref_img_len):cap_seq_len + sum(ref_img_len) + img_len]
+
+        return (
+            cap_freqs_cis,
+            ref_img_freqs_cis,
+            img_freqs_cis,
+            freqs_cis,
+            l_effective_cap_len,
+            seq_lengths,
+        )
diff --git a/modules/omnigen2/models/transformers/transformer_omnigen2.py b/modules/omnigen2/models/transformers/transformer_omnigen2.py
new file mode 100644
index 000000000..fe65b3170
--- /dev/null
+++ b/modules/omnigen2/models/transformers/transformer_omnigen2.py
@@ -0,0 +1,608 @@
+import itertools
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from torch.nn import RMSNorm
+from einops import rearrange
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.models.attention_processor import Attention
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import LuminaLayerNormContinuous, LuminaRMSNormZero
+from diffusers.models.attention import LuminaFeedForward
+
+from .block_lumina2 import Lumina2CombinedTimestepCaptionEmbedding
+from ..attention_processor import OmniGen2AttnProcessor
+from .repo import OmniGen2RotaryPosEmbed
+
+logger = logging.get_logger(__name__)
+
+
+class OmniGen2TransformerBlock(nn.Module):
+    """
+    Transformer block for OmniGen2 model.
+
+    This block implements a transformer layer with:
+    - Multi-head attention with flash attention
+    - Feed-forward network with SwiGLU activation
+    - RMS normalization
+    - Optional modulation for conditional generation
+
+    Args:
+        dim: Dimension of the input and output tensors
+        num_attention_heads: Number of attention heads
+        num_kv_heads: Number of key-value heads
+        multiple_of: Multiple of which the hidden dimension should be
+        ffn_dim_multiplier: Multiplier for the feed-forward network dimension
+        norm_eps: Epsilon value for normalization layers
+        modulation: Whether to use modulation for conditional generation
+        use_fused_rms_norm: Whether to use fused RMS normalization
+        use_fused_swiglu: Whether to use fused SwiGLU activation
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        num_kv_heads: int,
+        multiple_of: int,
+        ffn_dim_multiplier: float,
+        norm_eps: float,
+        modulation: bool = True,
+    ) -> None:
+        """Initialize the transformer block."""
+        super().__init__()
+        self.head_dim = dim // num_attention_heads
+        self.modulation = modulation
+
+        processor = OmniGen2AttnProcessor()
+        # Initialize attention layer
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=dim // num_attention_heads,
+            qk_norm="rms_norm",
+            heads=num_attention_heads,
+            kv_heads=num_kv_heads,
+            eps=1e-5,
+            bias=False,
+            out_bias=False,
+            processor=processor,
+        )
+
+        # Initialize feed-forward network
+        self.feed_forward = LuminaFeedForward(
+            dim=dim,
+            inner_dim=4 * dim,
+            multiple_of=multiple_of,
+            ffn_dim_multiplier=ffn_dim_multiplier
+        )
+
+        # Initialize normalization layers
+        if modulation:
+            self.norm1 = LuminaRMSNormZero(
+                embedding_dim=dim,
+                norm_eps=norm_eps,
+                norm_elementwise_affine=True
+            )
+        else:
+            self.norm1 = RMSNorm(dim, eps=norm_eps)
+
+        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+        self.initialize_weights()
+
+    def initialize_weights(self) -> None:
+        """
+        Initialize the weights of the transformer block.
+
+        Uses Xavier uniform initialization for linear layers and zero initialization for biases.
+        """
+        nn.init.xavier_uniform_(self.attn.to_q.weight)
+        nn.init.xavier_uniform_(self.attn.to_k.weight)
+        nn.init.xavier_uniform_(self.attn.to_v.weight)
+        nn.init.xavier_uniform_(self.attn.to_out[0].weight)
+
+        nn.init.xavier_uniform_(self.feed_forward.linear_1.weight)
+        nn.init.xavier_uniform_(self.feed_forward.linear_2.weight)
+        nn.init.xavier_uniform_(self.feed_forward.linear_3.weight)
+
+        if self.modulation:
+            nn.init.zeros_(self.norm1.linear.weight)
+            nn.init.zeros_(self.norm1.linear.bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        image_rotary_emb: torch.Tensor,
+        temb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Forward pass of the transformer block.
+
+        Args:
+            hidden_states: Input hidden states tensor
+            attention_mask: Attention mask tensor
+            image_rotary_emb: Rotary embeddings for image tokens
+            temb: Optional timestep embedding tensor
+
+        Returns:
+            torch.Tensor: Output hidden states after transformer block processing
+        """
+        import time
+        if self.modulation:
+            if temb is None:
+                raise ValueError("temb must be provided when modulation is enabled")
+
+            norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
+            attn_output = self.attn(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_hidden_states,
+                attention_mask=attention_mask,
+                image_rotary_emb=image_rotary_emb,
+            )
+            hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output)
+            mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1)))
+            hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output)
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+            attn_output = self.attn(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_hidden_states,
+                attention_mask=attention_mask,
+                image_rotary_emb=image_rotary_emb,
+            )
+            hidden_states = hidden_states + self.norm2(attn_output)
+            mlp_output = self.feed_forward(self.ffn_norm1(hidden_states))
+            hidden_states = hidden_states + self.ffn_norm2(mlp_output)
+
+        return hidden_states
+
+
+class OmniGen2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    """
+    OmniGen2 Transformer 2D Model.
+
+    A transformer-based diffusion model for image generation with:
+    - Patch-based image processing
+    - Rotary position embeddings
+    - Multi-head attention
+    - Conditional generation support
+
+    Args:
+        patch_size: Size of image patches
+        in_channels: Number of input channels
+        out_channels: Number of output channels (defaults to in_channels)
+        hidden_size: Size of hidden layers
+        num_layers: Number of transformer layers
+        num_refiner_layers: Number of refiner layers
+        num_attention_heads: Number of attention heads
+        num_kv_heads: Number of key-value heads
+        multiple_of: Multiple of which the hidden dimension should be
+        ffn_dim_multiplier: Multiplier for feed-forward network dimension
+        norm_eps: Epsilon value for normalization layers
+        axes_dim_rope: Dimensions for rotary position embeddings
+        axes_lens: Lengths for rotary position embeddings
+        text_feat_dim: Dimension of text features
+        timestep_scale: Scale factor for timestep embeddings
+        use_fused_rms_norm: Whether to use fused RMS normalization
+        use_fused_swiglu: Whether to use fused SwiGLU activation
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["Omnigen2TransformerBlock"]
+    _skip_layerwise_casting_patterns = ["x_embedder", "norm"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        out_channels: Optional[int] = None,
+        hidden_size: int = 2304,
+        num_layers: int = 26,
+        num_refiner_layers: int = 2,
+        num_attention_heads: int = 24,
+        num_kv_heads: int = 8,
+        multiple_of: int = 256,
+        ffn_dim_multiplier: Optional[float] = None,
+        norm_eps: float = 1e-5,
+        axes_dim_rope: Tuple[int, int, int] = (32, 32, 32),
+        axes_lens: Tuple[int, int, int] = (300, 512, 512),
+        text_feat_dim: int = 1024,
+        timestep_scale: float = 1.0
+    ) -> None:
+        """Initialize the OmniGen2 transformer model."""
+        super().__init__()
+
+        # Validate configuration
+        if (hidden_size // num_attention_heads) != sum(axes_dim_rope):
+            raise ValueError(
+                f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) "
+                f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})"
+            )
+
+        self.out_channels = out_channels or in_channels
+
+        # Initialize embeddings
+        self.rope_embedder = OmniGen2RotaryPosEmbed(
+            theta=10000,
+            axes_dim=axes_dim_rope,
+            axes_lens=axes_lens,
+            patch_size=patch_size,
+        )
+
+        self.x_embedder = nn.Linear(
+            in_features=patch_size * patch_size * in_channels,
+            out_features=hidden_size,
+        )
+
+        self.ref_image_patch_embedder = nn.Linear(
+            in_features=patch_size * patch_size * in_channels,
+            out_features=hidden_size,
+        )
+
+        self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding(
+            hidden_size=hidden_size,
+            text_feat_dim=text_feat_dim,
+            norm_eps=norm_eps,
+            timestep_scale=timestep_scale
+        )
+
+        # Initialize transformer blocks
+        self.noise_refiner = nn.ModuleList([
+            OmniGen2TransformerBlock(
+                hidden_size,
+                num_attention_heads,
+                num_kv_heads,
+                multiple_of,
+                ffn_dim_multiplier,
+                norm_eps,
+                modulation=True
+            )
+            for _ in range(num_refiner_layers)
+        ])
+
+        self.ref_image_refiner = nn.ModuleList([
+            OmniGen2TransformerBlock(
+                hidden_size,
+                num_attention_heads,
+                num_kv_heads,
+                multiple_of,
+                ffn_dim_multiplier,
+                norm_eps,
+                modulation=True
+            )
+            for _ in range(num_refiner_layers)
+        ])
+
+        self.context_refiner = nn.ModuleList(
+            [
+                OmniGen2TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=False
+                )
+                for _ in range(num_refiner_layers)
+            ]
+        )
+
+        # 3. Transformer blocks
+        self.layers = nn.ModuleList(
+            [
+                OmniGen2TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=True
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 4. Output norm & projection
+        self.norm_out = LuminaLayerNormContinuous(
+            embedding_dim=hidden_size,
+            conditioning_embedding_dim=min(hidden_size, 1024),
+            elementwise_affine=False,
+            eps=1e-6,
+            bias=True,
+            out_dim=patch_size * patch_size * self.out_channels
+        )
+
+        # Add learnable embeddings to distinguish different images
+        self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) # support max 5 ref images
+
+        self.gradient_checkpointing = False
+
+        self.initialize_weights()
+
+    def initialize_weights(self) -> None:
+        """
+        Initialize the weights of the model.
+
+        Uses Xavier uniform initialization for linear layers.
+        """
+        nn.init.xavier_uniform_(self.x_embedder.weight)
+        nn.init.constant_(self.x_embedder.bias, 0.0)
+
+        nn.init.xavier_uniform_(self.ref_image_patch_embedder.weight)
+        nn.init.constant_(self.ref_image_patch_embedder.bias, 0.0)
+
+        nn.init.zeros_(self.norm_out.linear_1.weight)
+        nn.init.zeros_(self.norm_out.linear_1.bias)
+        nn.init.zeros_(self.norm_out.linear_2.weight)
+        nn.init.zeros_(self.norm_out.linear_2.bias)
+
+        nn.init.normal_(self.image_index_embedding, std=0.02)
+
+    def img_patch_embed_and_refine(
+        self,
+        hidden_states,
+        ref_image_hidden_states,
+        padded_img_mask,
+        padded_ref_img_mask,
+        noise_rotary_emb,
+        ref_img_rotary_emb,
+        l_effective_ref_img_len,
+        l_effective_img_len,
+        temb
+    ):
+        batch_size = len(hidden_states)
+        max_combined_img_len = max([img_len + sum(ref_img_len) for img_len, ref_img_len in zip(l_effective_img_len, l_effective_ref_img_len)])
+
+        hidden_states = self.x_embedder(hidden_states)
+        ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states)
+
+        for i in range(batch_size):
+            shift = 0
+            for j, ref_img_len in enumerate(l_effective_ref_img_len[i]):
+                ref_image_hidden_states[i, shift:shift + ref_img_len, :] = ref_image_hidden_states[i, shift:shift + ref_img_len, :] + self.image_index_embedding[j]
+                shift += ref_img_len
+
+        for layer in self.noise_refiner:
+            hidden_states = layer(hidden_states, padded_img_mask, noise_rotary_emb, temb)
+
+        flat_l_effective_ref_img_len = list(itertools.chain(*l_effective_ref_img_len))
+        num_ref_images = len(flat_l_effective_ref_img_len)
+        max_ref_img_len = max(flat_l_effective_ref_img_len)
+
+        batch_ref_img_mask = ref_image_hidden_states.new_zeros(num_ref_images, max_ref_img_len, dtype=torch.bool)
+        batch_ref_image_hidden_states = ref_image_hidden_states.new_zeros(num_ref_images, max_ref_img_len, self.config.hidden_size)
+        batch_ref_img_rotary_emb = hidden_states.new_zeros(num_ref_images, max_ref_img_len, ref_img_rotary_emb.shape[-1], dtype=ref_img_rotary_emb.dtype)
+        batch_temb = temb.new_zeros(num_ref_images, *temb.shape[1:], dtype=temb.dtype)
+
+        # sequence of ref imgs to batch
+        idx = 0
+        for i in range(batch_size):
+            shift = 0
+            for ref_img_len in l_effective_ref_img_len[i]:
+                batch_ref_img_mask[idx, :ref_img_len] = True
+                batch_ref_image_hidden_states[idx, :ref_img_len] = ref_image_hidden_states[i, shift:shift + ref_img_len]
+                batch_ref_img_rotary_emb[idx, :ref_img_len] = ref_img_rotary_emb[i, shift:shift + ref_img_len]
+                batch_temb[idx] = temb[i]
+                shift += ref_img_len
+                idx += 1
+
+        # refine ref imgs separately
+        for layer in self.ref_image_refiner:
+            batch_ref_image_hidden_states = layer(batch_ref_image_hidden_states, batch_ref_img_mask, batch_ref_img_rotary_emb, batch_temb)
+
+        # batch of ref imgs to sequence
+        idx = 0
+        for i in range(batch_size):
+            shift = 0
+            for ref_img_len in l_effective_ref_img_len[i]:
+                ref_image_hidden_states[i, shift:shift + ref_img_len] = batch_ref_image_hidden_states[idx, :ref_img_len]
+                shift += ref_img_len
+                idx += 1
+
+        combined_img_hidden_states = hidden_states.new_zeros(batch_size, max_combined_img_len, self.config.hidden_size)
+        for i, (ref_img_len, img_len) in enumerate(zip(l_effective_ref_img_len, l_effective_img_len)):
+            combined_img_hidden_states[i, :sum(ref_img_len)] = ref_image_hidden_states[i, :sum(ref_img_len)]
+            combined_img_hidden_states[i, sum(ref_img_len):sum(ref_img_len) + img_len] = hidden_states[i, :img_len]
+
+        return combined_img_hidden_states
+
+    def flat_and_pad_to_seq(self, hidden_states, ref_image_hidden_states):
+        batch_size = len(hidden_states)
+        p = self.config.patch_size
+        device = hidden_states[0].device
+
+        img_sizes = [(img.size(1), img.size(2)) for img in hidden_states]
+        l_effective_img_len = [(H // p) * (W // p) for (H, W) in img_sizes]
+
+        if ref_image_hidden_states is not None:
+            ref_img_sizes = [[(img.size(1), img.size(2)) for img in imgs] if imgs is not None else None for imgs in ref_image_hidden_states]
+            l_effective_ref_img_len = [[(ref_img_size[0] // p) * (ref_img_size[1] // p) for ref_img_size in _ref_img_sizes] if _ref_img_sizes is not None else [0] for _ref_img_sizes in ref_img_sizes]
+        else:
+            ref_img_sizes = [None for _ in range(batch_size)]
+            l_effective_ref_img_len = [[0] for _ in range(batch_size)]
+
+        max_ref_img_len = max([sum(ref_img_len) for ref_img_len in l_effective_ref_img_len])
+        max_img_len = max(l_effective_img_len)
+
+        # ref image patch embeddings
+        flat_ref_img_hidden_states = []
+        for i in range(batch_size):
+            if ref_img_sizes[i] is not None:
+                imgs = []
+                for ref_img in ref_image_hidden_states[i]:
+                    C, H, W = ref_img.size()
+                    ref_img = rearrange(ref_img, 'c (h p1) (w p2) -> (h w) (p1 p2 c)', p1=p, p2=p)
+                    imgs.append(ref_img)
+
+                img = torch.cat(imgs, dim=0)
+                flat_ref_img_hidden_states.append(img)
+            else:
+                flat_ref_img_hidden_states.append(None)
+
+        # image patch embeddings
+        flat_hidden_states = []
+        for i in range(batch_size):
+            img = hidden_states[i]
+            C, H, W = img.size()
+
+            img = rearrange(img, 'c (h p1) (w p2) -> (h w) (p1 p2 c)', p1=p, p2=p)
+            flat_hidden_states.append(img)
+
+        padded_ref_img_hidden_states = torch.zeros(batch_size, max_ref_img_len, flat_hidden_states[0].shape[-1], device=device, dtype=flat_hidden_states[0].dtype)
+        padded_ref_img_mask = torch.zeros(batch_size, max_ref_img_len, dtype=torch.bool, device=device)
+        for i in range(batch_size):
+            if ref_img_sizes[i] is not None:
+                padded_ref_img_hidden_states[i, :sum(l_effective_ref_img_len[i])] = flat_ref_img_hidden_states[i]
+                padded_ref_img_mask[i, :sum(l_effective_ref_img_len[i])] = True
+
+        padded_hidden_states = torch.zeros(batch_size, max_img_len, flat_hidden_states[0].shape[-1], device=device, dtype=flat_hidden_states[0].dtype)
+        padded_img_mask = torch.zeros(batch_size, max_img_len, dtype=torch.bool, device=device)
+        for i in range(batch_size):
+            padded_hidden_states[i, :l_effective_img_len[i]] = flat_hidden_states[i]
+            padded_img_mask[i, :l_effective_img_len[i]] = True
+
+        return (
+            padded_hidden_states,
+            padded_ref_img_hidden_states,
+            padded_img_mask,
+            padded_ref_img_mask,
+            l_effective_ref_img_len,
+            l_effective_img_len,
+            ref_img_sizes,
+            img_sizes,
+        )
+
+    def forward(
+        self,
+        hidden_states: Union[torch.Tensor, List[torch.Tensor]],
+        timestep: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        ref_image_hidden_states: Optional[List[List[torch.Tensor]]] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = False,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = 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 attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        # 1. Condition, positional & patch embedding
+        batch_size = len(hidden_states)
+        is_hidden_states_tensor = isinstance(hidden_states, torch.Tensor)
+
+        if is_hidden_states_tensor:
+            assert hidden_states.ndim == 4
+            hidden_states = [_hidden_states for _hidden_states in hidden_states]
+
+        device = hidden_states[0].device
+
+        temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states[0].dtype)
+
+        (
+            hidden_states,
+            ref_image_hidden_states,
+            img_mask,
+            ref_img_mask,
+            l_effective_ref_img_len,
+            l_effective_img_len,
+            ref_img_sizes,
+            img_sizes,
+        ) = self.flat_and_pad_to_seq(hidden_states, ref_image_hidden_states)
+
+        (
+            context_rotary_emb,
+            ref_img_rotary_emb,
+            noise_rotary_emb,
+            rotary_emb,
+            encoder_seq_lengths,
+            seq_lengths,
+        ) = self.rope_embedder(
+            freqs_cis,
+            text_attention_mask,
+            l_effective_ref_img_len,
+            l_effective_img_len,
+            ref_img_sizes,
+            img_sizes,
+            device,
+        )
+
+        # 2. Context refinement
+        for layer in self.context_refiner:
+            text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb)
+
+        combined_img_hidden_states = self.img_patch_embed_and_refine(
+            hidden_states,
+            ref_image_hidden_states,
+            img_mask,
+            ref_img_mask,
+            noise_rotary_emb,
+            ref_img_rotary_emb,
+            l_effective_ref_img_len,
+            l_effective_img_len,
+            temb,
+        )
+
+        # 3. Joint Transformer blocks
+        max_seq_len = max(seq_lengths)
+
+        attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
+        joint_hidden_states = hidden_states.new_zeros(batch_size, max_seq_len, self.config.hidden_size)
+        for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+            attention_mask[i, :seq_len] = True
+            joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len]
+            joint_hidden_states[i, encoder_seq_len:seq_len] = combined_img_hidden_states[i, :seq_len - encoder_seq_len]
+
+        hidden_states = joint_hidden_states
+
+        for layer_idx, layer in enumerate(self.layers):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    layer, hidden_states, attention_mask, rotary_emb, temb
+                )
+            else:
+                hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
+
+        # 4. Output norm & projection
+        hidden_states = self.norm_out(hidden_states, temb)
+
+        p = self.config.patch_size
+        output = []
+        for i, (img_size, img_len, seq_len) in enumerate(zip(img_sizes, l_effective_img_len, seq_lengths)):
+            height, width = img_size
+            output.append(rearrange(hidden_states[i][seq_len - img_len:seq_len], '(h w) (p1 p2 c) -> c (h p1) (w p2)', h=height // p, w=width // p, p1=p, p2=p))
+        if is_hidden_states_tensor:
+            output = torch.stack(output, dim=0)
+
+        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)
diff --git a/modules/omnigen2/pipeline_omnigen2.py b/modules/omnigen2/pipeline_omnigen2.py
new file mode 100644
index 000000000..a7d6c2ca3
--- /dev/null
+++ b/modules/omnigen2/pipeline_omnigen2.py
@@ -0,0 +1,718 @@
+"""
+OmniGen2 Diffusion Pipeline
+
+Copyright 2025 BAAI, The OmniGen2 Team and 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 typing import Any, Dict, List, Optional, Tuple, Union
+from dataclasses import dataclass
+import inspect
+import numpy as np
+import torch
+import torch.nn.functional as F
+import PIL.Image
+
+from transformers import Qwen2_5_VLForConditionalGeneration
+from diffusers.utils import BaseOutput
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    is_torch_xla_available,
+    logging,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+from .models.transformers import OmniGen2Transformer2DModel
+from .models.transformers.repo import OmniGen2RotaryPosEmbed
+from .image_processor import OmniGen2ImageProcessor
+
+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
+
+@dataclass
+class FMPipelineOutput(BaseOutput):
+    """
+    Output class for OmniGen2 pipeline.
+
+    Args:
+        images (Union[List[PIL.Image.Image], np.ndarray]): 
+            List of denoised PIL images of length `batch_size` or numpy array of shape 
+            `(batch_size, height, width, num_channels)`. Contains the generated images.
+    """
+    images: Union[List[PIL.Image.Image], np.ndarray]
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class OmniGen2Pipeline(DiffusionPipeline):
+    """
+    Pipeline for text-to-image generation using OmniGen2.
+
+    This pipeline implements a text-to-image generation model that uses:
+    - Qwen2.5-VL for text encoding
+    - A custom transformer architecture for image generation
+    - VAE for image encoding/decoding
+    - FlowMatchEulerDiscreteScheduler for noise scheduling
+
+    Args:
+        transformer (OmniGen2Transformer2DModel): The transformer model for image generation.
+        vae (AutoencoderKL): The VAE model for image encoding/decoding.
+        scheduler (FlowMatchEulerDiscreteScheduler): The scheduler for noise scheduling.
+        text_encoder (Qwen2_5_VLModel): The text encoder model.
+        tokenizer (Union[Qwen2Tokenizer, Qwen2TokenizerFast]): The tokenizer for text processing.
+    """
+
+    model_cpu_offload_seq = "mllm->transformer->vae"
+
+    def __init__(
+        self,
+        transformer: OmniGen2Transformer2DModel,
+        vae: AutoencoderKL,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        mllm: Qwen2_5_VLForConditionalGeneration,
+        processor,
+    ) -> None:
+        """
+        Initialize the OmniGen2 pipeline.
+
+        Args:
+            transformer: The transformer model for image generation.
+            vae: The VAE model for image encoding/decoding.
+            scheduler: The scheduler for noise scheduling.
+            text_encoder: The text encoder model.
+            tokenizer: The tokenizer for text processing.
+        """
+        super().__init__()
+
+        self.register_modules(
+            transformer=transformer,
+            vae=vae,
+            scheduler=scheduler,
+            mllm=mllm,
+            processor=processor
+        )
+        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 = OmniGen2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_resize=True)
+        self.default_sample_size = 128
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: Optional[torch.Generator],
+        latents: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        """
+        Prepare the initial latents for the diffusion process.
+
+        Args:
+            batch_size: The number of images to generate.
+            num_channels_latents: The number of channels in the latent space.
+            height: The height of the generated image.
+            width: The width of the generated image.
+            dtype: The data type of the latents.
+            device: The device to place the latents on.
+            generator: The random number generator to use.
+            latents: Optional pre-computed latents to use instead of random initialization.
+
+        Returns:
+            torch.FloatTensor: The prepared latents tensor.
+        """
+        height = int(height) // self.vae_scale_factor
+        width = int(width) // self.vae_scale_factor
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+        return latents
+
+    def encode_vae(self, img: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        Encode an image into the VAE latent space.
+
+        Args:
+            img: The input image tensor to encode.
+
+        Returns:
+            torch.FloatTensor: The encoded latent representation.
+        """
+        z0 = self.vae.encode(img.to(dtype=self.vae.dtype)).latent_dist.sample()
+        if self.vae.config.shift_factor is not None:
+            z0 = z0 - self.vae.config.shift_factor
+        if self.vae.config.scaling_factor is not None:
+            z0 = z0 * self.vae.config.scaling_factor
+        z0 = z0.to(dtype=self.vae.dtype)
+        return z0
+
+    def prepare_image(
+        self,
+        images: Union[List[PIL.Image.Image], PIL.Image.Image],
+        batch_size: int,
+        num_images_per_prompt: int,
+        max_pixels: int,
+        max_side_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> List[Optional[torch.FloatTensor]]:
+        """
+        Prepare input images for processing by encoding them into the VAE latent space.
+
+        Args:
+            images: Single image or list of images to process.
+            batch_size: The number of images to generate per prompt.
+            num_images_per_prompt: The number of images to generate for each prompt.
+            device: The device to place the encoded latents on.
+            dtype: The data type of the encoded latents.
+
+        Returns:
+            List[Optional[torch.FloatTensor]]: List of encoded latent representations for each image.
+        """
+        if batch_size == 1:
+            images = [images]
+        latents = []
+        for i, img in enumerate(images):
+            if img is not None and len(img) > 0:
+                ref_latents = []
+                for j, img_j in enumerate(img):
+                    img_j = self.image_processor.preprocess(img_j, max_pixels=max_pixels, max_side_length=max_side_length)
+                    ref_latents.append(self.encode_vae(img_j.to(device=device)).squeeze(0))
+            else:
+                ref_latents = None
+            for _ in range(num_images_per_prompt):
+                latents.append(ref_latents)
+
+        return latents
+
+    def _get_qwen2_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        max_sequence_length: int = 256,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Get prompt embeddings from the Qwen2 text encoder.
+
+        Args:
+            prompt: The prompt or list of prompts to encode.
+            device: The device to place the embeddings on. If None, uses the pipeline's device.
+            max_sequence_length: Maximum sequence length for tokenization.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+                - The prompt embeddings tensor
+                - The attention mask tensor
+
+        Raises:
+            Warning: If the input text is truncated due to sequence length limitations.
+        """
+        device = device or self._execution_device
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        # text_inputs = self.processor.tokenizer(
+        #     prompt,
+        #     padding="max_length",
+        #     max_length=max_sequence_length,
+        #     truncation=True,
+        #     return_tensors="pt",
+        # )
+        text_inputs = self.processor.tokenizer(
+            prompt,
+            padding="longest",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids.to(device)
+        untruncated_ids = self.processor.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(device)
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.processor.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because Gemma can only handle sequences up to"
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_attention_mask = text_inputs.attention_mask.to(device)
+        prompt_embeds = self.mllm(
+            text_input_ids,
+            attention_mask=prompt_attention_mask,
+            output_hidden_states=True,
+        ).hidden_states[-1]
+
+        if self.mllm is not None:
+            dtype = self.mllm.dtype
+        elif self.transformer is not None:
+            dtype = self.transformer.dtype
+        else:
+            dtype = None
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        return prompt_embeds, prompt_attention_mask
+
+    def _apply_chat_template(self, prompt: str):
+        prompt = [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant that generates high-quality images based on user instructions.",
+            },
+            {"role": "user", "content": prompt},
+        ]
+        prompt = self.processor.tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=False)
+        return prompt
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_attention_mask: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 256,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt 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`). For
+                Lumina-T2I, this should be "".
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                whether to use classifier free guidance or not
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                number of images that should be generated per prompt
+            device: (`torch.device`, *optional*):
+                torch device to place the resulting embeddings on
+            prompt_embeds (`torch.Tensor`, *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.Tensor`, *optional*):
+                Pre-generated negative text embeddings. For Lumina-T2I, it's should be the embeddings of the "" string.
+            max_sequence_length (`int`, defaults to `256`):
+                Maximum sequence length to use for the prompt.
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt = [self._apply_chat_template(_prompt) for _prompt in prompt]
+
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        if prompt_embeds is None:
+            prompt_embeds, prompt_attention_mask = self._get_qwen2_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                max_sequence_length=max_sequence_length
+            )
+
+        batch_size, 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)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+        prompt_attention_mask = prompt_attention_mask.view(batch_size * num_images_per_prompt, -1)
+
+        # Get negative embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt if negative_prompt is not None else ""
+
+            # Normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt = [self._apply_chat_template(_negative_prompt) for _negative_prompt in negative_prompt]
+
+            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 isinstance(negative_prompt, str):
+                negative_prompt = [negative_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_embeds, negative_prompt_attention_mask = self._get_qwen2_prompt_embeds(
+                prompt=negative_prompt,
+                device=device,
+                max_sequence_length=max_sequence_length,
+            )
+
+            batch_size, seq_len, _ = negative_prompt_embeds.shape
+            # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.view(
+                batch_size * num_images_per_prompt, -1
+            )
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def text_guidance_scale(self):
+        return self._text_guidance_scale
+
+    @property
+    def image_guidance_scale(self):
+        return self._image_guidance_scale
+
+    @property
+    def cfg_range(self):
+        return self._cfg_range
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.LongTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.LongTensor] = None,
+        max_sequence_length: Optional[int] = None,
+        callback_on_step_end_tensor_inputs: Optional[List[str]] = None,
+        input_images: Optional[List[PIL.Image.Image]] = None,
+        num_images_per_prompt: int = 1,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        max_pixels: int = 2048 * 2048,
+        max_input_image_side_length: int = 2048,
+        align_res: bool = True,
+        num_inference_steps: int = 28,
+        text_guidance_scale: float = 4.0,
+        image_guidance_scale: float = 1.0,
+        cfg_range: Tuple[float, float] = (0.0, 1.0),
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        timesteps: List[int] = None,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        verbose: bool = False,
+        step_func=None,
+    ):
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        self._text_guidance_scale = text_guidance_scale
+        self._image_guidance_scale = image_guidance_scale
+        self._cfg_range = cfg_range
+        self._attention_kwargs = attention_kwargs
+
+        # 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
+
+        # 3. Encode input prompt
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt,
+            self.text_guidance_scale > 1.0,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            max_sequence_length=max_sequence_length,
+        )
+
+        dtype = self.vae.dtype
+        # 3. Prepare control image
+        ref_latents = self.prepare_image(
+            images=input_images,
+            batch_size=batch_size,
+            num_images_per_prompt=num_images_per_prompt,
+            max_pixels=max_pixels,
+            max_side_length=max_input_image_side_length,
+            device=device,
+            dtype=dtype,
+        )
+
+        if input_images is None:
+            input_images = []
+
+        if len(input_images) == 1 and align_res:
+            width, height = ref_latents[0][0].shape[-1] * self.vae_scale_factor, ref_latents[0][0].shape[-2] * self.vae_scale_factor
+            ori_width, ori_height = width, height
+        else:
+            ori_width, ori_height = width, height
+
+            cur_pixels = height * width
+            ratio = (max_pixels / cur_pixels) ** 0.5
+            ratio = min(ratio, 1.0)
+
+            height, width = int(height * ratio) // 16 * 16, int(width * ratio) // 16 * 16
+
+        if len(input_images) == 0:
+            self._image_guidance_scale = 1
+
+        # 4. Prepare latents.
+        latent_channels = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            latent_channels,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        freqs_cis = OmniGen2RotaryPosEmbed.get_freqs_cis(
+            self.transformer.config.axes_dim_rope,
+            self.transformer.config.axes_lens,
+            theta=10000,
+        )
+
+        image = self.processing(
+            latents=latents,
+            ref_latents=ref_latents,
+            prompt_embeds=prompt_embeds,
+            freqs_cis=freqs_cis,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            device=device,
+            dtype=dtype,
+            verbose=verbose,
+            step_func=step_func,
+        )
+
+        image = F.interpolate(image, size=(ori_height, ori_width), mode='bilinear')
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return image
+        else:
+            return FMPipelineOutput(images=image)
+
+    def processing(
+        self,
+        latents,
+        ref_latents,
+        prompt_embeds,
+        freqs_cis,
+        negative_prompt_embeds,
+        prompt_attention_mask,
+        negative_prompt_attention_mask,
+        num_inference_steps,
+        timesteps,
+        device,
+        dtype,
+        verbose,
+        step_func=None
+    ):
+        batch_size = latents.shape[0]
+
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            num_tokens=latents.shape[-2] * latents.shape[-1]
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                model_pred = self.predict(
+                    t=t,
+                    latents=latents,
+                    prompt_embeds=prompt_embeds,
+                    freqs_cis=freqs_cis,
+                    prompt_attention_mask=prompt_attention_mask,
+                    ref_image_hidden_states=ref_latents,
+                )
+                text_guidance_scale = self.text_guidance_scale if self.cfg_range[0] <= i / len(timesteps) <= self.cfg_range[1] else 1.0
+                image_guidance_scale = self.image_guidance_scale if self.cfg_range[0] <= i / len(timesteps) <= self.cfg_range[1] else 1.0
+
+                if text_guidance_scale > 1.0 and image_guidance_scale > 1.0:
+                    model_pred_ref = self.predict(
+                        t=t,
+                        latents=latents,
+                        prompt_embeds=negative_prompt_embeds,
+                        freqs_cis=freqs_cis,
+                        prompt_attention_mask=negative_prompt_attention_mask,
+                        ref_image_hidden_states=ref_latents,
+                    )
+
+                    if image_guidance_scale != 1:
+                        model_pred_uncond = self.predict(
+                            t=t,
+                            latents=latents,
+                            prompt_embeds=negative_prompt_embeds,
+                            freqs_cis=freqs_cis,
+                            prompt_attention_mask=negative_prompt_attention_mask,
+                            ref_image_hidden_states=None,
+                        )
+                    else:
+                        model_pred_uncond = torch.zeros_like(model_pred)
+
+                    model_pred = model_pred_uncond + image_guidance_scale * (model_pred_ref - model_pred_uncond) + \
+                    text_guidance_scale * (model_pred - model_pred_ref)
+                elif text_guidance_scale > 1.0:
+                    model_pred_uncond = self.predict(
+                        t=t,
+                        latents=latents,
+                        prompt_embeds=negative_prompt_embeds,
+                        freqs_cis=freqs_cis,
+                        prompt_attention_mask=negative_prompt_attention_mask,
+                        ref_image_hidden_states=None,
+                    )
+                    model_pred = model_pred_uncond + text_guidance_scale * (model_pred - model_pred_uncond)
+
+                latents = self.scheduler.step(model_pred, t, latents, return_dict=False)[0]
+
+                latents = latents.to(dtype=dtype)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if step_func is not None:
+                    step_func(i, self._num_timesteps)
+
+        latents = latents.to(dtype=dtype)
+        if self.vae.config.scaling_factor is not None:
+            latents = latents / self.vae.config.scaling_factor
+        if self.vae.config.shift_factor is not None:
+            latents = latents + self.vae.config.shift_factor
+        image = self.vae.decode(latents, return_dict=False)[0]
+
+        return image
+
+    def predict(
+        self,
+        t,
+        latents,
+        prompt_embeds,
+        freqs_cis,
+        prompt_attention_mask,
+        ref_image_hidden_states,
+    ):
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+        batch_size, num_channels_latents, height, width = latents.shape
+
+        optional_kwargs = {}
+        if 'ref_image_hidden_states' in set(inspect.signature(self.transformer.forward).parameters.keys()):
+            optional_kwargs['ref_image_hidden_states'] = ref_image_hidden_states
+
+        model_pred = self.transformer(
+            latents,
+            timestep,
+            prompt_embeds,
+            freqs_cis,
+            prompt_attention_mask,
+            **optional_kwargs
+        )
+        return model_pred
diff --git a/modules/paths.py b/modules/paths.py
index 06cf34bff..d3257f043 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -29,6 +29,7 @@ script_path = os.path.dirname(modules_path)
 data_path = cli.data_dir
 models_config = cli.models_dir or config.get('models_dir') or 'models'
 models_path = models_config if os.path.isabs(models_config) else os.path.join(data_path, models_config)
+params_path = os.environ.get('SD_PATH_PARAMS', os.path.join(data_path, "params.txt"))
 extensions_dir = cli.extensions_dir or os.path.join(data_path, "extensions")
 extensions_builtin_dir = "extensions-builtin"
 sd_configs_path = os.path.join(script_path, "configs")
diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py
index d3b148805..67dbb7b4d 100644
--- a/modules/postprocess/yolo.py
+++ b/modules/postprocess/yolo.py
@@ -383,7 +383,7 @@ class YoloRestorer(Detailer):
                 max_size = shared.opts.detailer_max_size if shared.opts.detailer_max_size < 1 and shared.opts.detailer_max_size > 0 else 1.0
                 max_size = gr.Slider(label="Max size", elem_id=f"{tab}_detailer_max_size", value=max_size, minimum=0.0, maximum=1.0, step=0.05)
             with gr.Row(elem_classes=['flex-break']):
-                renoise_value = gr.Slider(minimum=0.5, maximum=1.5, step=0.01, label='Renoiose', value=shared.opts.detailer_sigma_adjust, elem_id=f"{tab}_detailer_renoise")
+                renoise_value = gr.Slider(minimum=0.5, maximum=1.5, step=0.01, label='Renoise', value=shared.opts.detailer_sigma_adjust, elem_id=f"{tab}_detailer_renoise")
                 renoise_end = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Renoise end', value=shared.opts.detailer_sigma_adjust_max, elem_id=f"{tab}_detailer_renoise_end")
             detailers.change(fn=ui_settings_change, inputs=[detailers, classes, strength, padding, blur, min_confidence, max_detected, min_size, max_size, iou, steps, renoise_value, renoise_end], outputs=[])
             classes.change(fn=ui_settings_change, inputs=[detailers, classes, strength, padding, blur, min_confidence, max_detected, min_size, max_size, iou, steps, renoise_value, renoise_end], outputs=[])
diff --git a/modules/processing_args.py b/modules/processing_args.py
index 6ef43667b..36f637873 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -48,17 +48,21 @@ def task_specific_kwargs(p, model):
             'image': p.init_images,
             'strength': p.denoising_strength,
         }
-        if model.__class__.__name__ == 'FluxImg2ImgPipeline': # needs explicit width/height
+        if model.__class__.__name__ == 'FluxImg2ImgPipeline' or model.__class__.__name__ == 'FluxKontextPipeline': # needs explicit width/height
             if torch.is_tensor(p.init_images[0]):
-                p.width = p.init_images[0].shape[-1] * 16
-                p.height = p.init_images[0].shape[-2] * 16
+                p.width, p.height = p.init_images[0].shape[-1] * 16, p.init_images[0].shape[-2] * 16
             else:
-                p.width = 8 * math.ceil(p.init_images[0].width / 8)
-                p.height = 8 * math.ceil(p.init_images[0].height / 8)
+                p.width, p.height = 8 * math.ceil(p.init_images[0].width / 8), 8 * math.ceil(p.init_images[0].height / 8)
+            if model.__class__.__name__ == 'FluxKontextPipeline':
+                aspect_ratio = p.width / p.height
+                vae_scale_factor = 16
+                max_area = max(p.width, p.height)**2
+                p.width, p.height = round((max_area * aspect_ratio) ** 0.5), round((max_area / aspect_ratio) ** 0.5)
+                p.width, p.height = p.width // vae_scale_factor * vae_scale_factor, p.height // vae_scale_factor * vae_scale_factor
+                task_args['max_area'] = max_area
             task_args['width'], task_args['height'] = p.width, p.height
-        if model.__class__.__name__ == 'OmniGenPipeline':
-            p.width = 16 * math.ceil(p.init_images[0].width / 16)
-            p.height = 16 * math.ceil(p.init_images[0].height / 16)
+        elif model.__class__.__name__ == 'OmniGenPipeline' or model.__class__.__name__ == 'OmniGen2Pipeline':
+            p.width, p.height = 16 * math.ceil(p.init_images[0].width / 16), 16 * math.ceil(p.init_images[0].height / 16)
             task_args = {
                 'width': p.width,
                 'height': p.height,
@@ -144,12 +148,15 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     steps = kwargs.get("num_inference_steps", None) or len(getattr(p, 'timesteps', ['1']))
     clip_skip = kwargs.pop("clip_skip", 1)
 
+    extra_networks.activate(p, include=['text_encoder', 'text_encoder_2', 'text_encoder_3'])
+
     parser = 'fixed'
     prompt_attention = prompt_attention or shared.opts.prompt_attention
     if (prompt_attention != 'fixed') and ('Onnx' not in model.__class__.__name__) and ('prompt' not in p.task_args) and (
         'StableDiffusion' in model.__class__.__name__ or
         'StableCascade' in model.__class__.__name__ or
         'Flux' in model.__class__.__name__ or
+        'Chroma' in model.__class__.__name__ or
         'HiDreamImagePipeline' in model.__class__.__name__ # hidream-e1 has different embeds
     ):
         try:
@@ -163,10 +170,9 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     else:
         prompt_parser_diffusers.embedder = None
 
-    extra_networks.activate(p, include=['text_encoder', 'text_encoder_2', 'text_encoder_3'])
     if 'prompt' in possible:
         if 'OmniGen' in model.__class__.__name__:
-            prompts = [p.replace('|image|', '<|image_1|>') for p in prompts]
+            prompts = [p.replace('|image|', '<img><|image_1|></img>') for p in prompts]
         if 'HiDreamImage' in model.__class__.__name__  and prompt_parser_diffusers.embedder is not None:
             args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
             prompt_embeds = prompt_parser_diffusers.embedder('prompt_embeds')
@@ -174,15 +180,16 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             args['prompt_embeds_llama3'] = prompt_embeds[1]
         elif hasattr(model, 'text_encoder') and hasattr(model, 'tokenizer') and 'prompt_embeds' in possible and prompt_parser_diffusers.embedder is not None:
             args['prompt_embeds'] = prompt_parser_diffusers.embedder('prompt_embeds')
-            if prompt_parser_diffusers.embedder is not None:
-                if 'StableCascade' in model.__class__.__name__:
-                    args['prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('positive_pooleds').unsqueeze(0)
-                elif 'XL' in model.__class__.__name__:
-                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
-                elif 'StableDiffusion3' in model.__class__.__name__:
-                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
-                elif 'Flux' in model.__class__.__name__:
-                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+            if 'StableCascade' in model.__class__.__name__:
+                args['prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('positive_pooleds').unsqueeze(0)
+            elif 'XL' in model.__class__.__name__:
+                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+            elif 'StableDiffusion3' in model.__class__.__name__:
+                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+            elif 'Flux' in model.__class__.__name__:
+                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+            elif 'Chroma' in model.__class__.__name__:
+                args['prompt_attention_mask'] = prompt_parser_diffusers.embedder('prompt_attention_masks')
         else:
             args['prompt'] = prompts
     if 'negative_prompt' in possible:
@@ -193,13 +200,14 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             args['negative_prompt_embeds_llama3'] = negative_prompt_embeds[1]
         elif hasattr(model, 'text_encoder') and hasattr(model, 'tokenizer') and 'negative_prompt_embeds' in possible and prompt_parser_diffusers.embedder is not None:
             args['negative_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_prompt_embeds')
-            if prompt_parser_diffusers.embedder is not None:
-                if 'StableCascade' in model.__class__.__name__:
-                    args['negative_prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('negative_pooleds').unsqueeze(0)
-                elif 'XL' in model.__class__.__name__:
-                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
-                elif 'StableDiffusion3' in model.__class__.__name__:
-                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+            if 'StableCascade' in model.__class__.__name__:
+                args['negative_prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('negative_pooleds').unsqueeze(0)
+            elif 'XL' in model.__class__.__name__:
+                args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+            elif 'StableDiffusion3' in model.__class__.__name__:
+                args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+            elif 'Chroma' in model.__class__.__name__:
+                args['negative_prompt_attention_mask'] = prompt_parser_diffusers.embedder('negative_prompt_attention_masks')
         else:
             if 'PixArtSigmaPipeline' in model.__class__.__name__: # pixart-sigma pipeline throws list-of-list for negative prompt
                 args['negative_prompt'] = negative_prompts[0]
@@ -278,7 +286,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             kwargs['output_type'] = 'np' # only set latent if model has vae
 
     # model specific
-    if 'Kandinsky' in model.__class__.__name__:
+    if 'Kandinsky' in model.__class__.__name__ or 'Cosmos2' in model.__class__.__name__ or 'OmniGen2' in model.__class__.__name__:
         kwargs['output_type'] = 'np' # only set latent if model has vae
     if 'StableCascade' in model.__class__.__name__:
         kwargs.pop("guidance_scale") # remove
diff --git a/modules/processing_class.py b/modules/processing_class.py
index 6b10f07a7..fc611b3ce 100644
--- a/modules/processing_class.py
+++ b/modules/processing_class.py
@@ -152,6 +152,8 @@ class StableDiffusionProcessing:
         self.positive_pooleds = []
         self.negative_embeds = []
         self.negative_pooleds = []
+        self.prompt_attention_masks = []
+        self.negative_prompt_attention_masks = []
         self.disable_extra_networks = False
         self.iteration = 0
         self.network_data = {}
@@ -256,7 +258,8 @@ class StableDiffusionProcessing:
         self.all_subseeds = None
 
         # a1111 compatibility items
-        shared.opts.data['clip_skip'] = int(self.clip_skip) # for compatibility with a1111 sd_hijack_clip
+        if not shared.native:
+            shared.opts.data['clip_skip'] = int(self.clip_skip) # for compatibility with a1111 sd_hijack_clip
         self.seed_enable_extras: bool = True
         self.is_using_inpainting_conditioning = False # a111 compatibility
         self.batch_index = 0
@@ -321,6 +324,8 @@ class StableDiffusionProcessing:
         self.positive_pooleds = []
         self.negative_embeds = []
         self.negative_pooleds = []
+        self.prompt_attention_masks = []
+        self.negative_prompt_attention_mask = []
 
     def __str__(self):
         return f'{self.__class__.__name__}: {self.__dict__}'
@@ -459,8 +464,8 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         if all_subseeds is not None:
             self.all_subseeds = all_subseeds
 
-        if self.sampler_name == "PLMS":
-            self.sampler_name = 'UniPC'
+        if self.sampler_name == 'PLMS':
+            self.sampler_name = 'Default'
         if not shared.native:
             self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)
             if hasattr(self.sampler, "initialize"):
diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py
index 03d0b7b78..4b10b4835 100644
--- a/modules/processing_diffusers.py
+++ b/modules/processing_diffusers.py
@@ -170,7 +170,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
         prev_job = shared.state.job
 
         # hires runs on original pipeline
-        if hasattr(shared.sd_model, 'restore_pipeline') and shared.sd_model.restore_pipeline is not None:
+        if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None) and not shared.opts.control_hires:
             shared.sd_model.restore_pipeline()
 
         # upscale
diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py
index 0f2d7bc6c..5d1647f8f 100644
--- a/modules/processing_helpers.py
+++ b/modules/processing_helpers.py
@@ -97,10 +97,10 @@ def get_sampler_name(sampler_index: int, img: bool = False) -> str:
     if len(sd_samplers.samplers) > sampler_index:
         sampler_name = sd_samplers.samplers[sampler_index].name
     else:
-        sampler_name = "UniPC"
+        sampler_name = "Default"
         shared.log.warning(f'Sampler not found: index={sampler_index} available={[s.name for s in sd_samplers.samplers]} fallback={sampler_name}')
     if img and sampler_name == "PLMS":
-        sampler_name = "UniPC"
+        sampler_name = "Default"
         shared.log.warning(f'Sampler not compatible: name=PLMS fallback={sampler_name}')
     return sampler_name
 
@@ -424,19 +424,28 @@ def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler
 def fix_prompts(p, prompts, negative_prompts, prompts_2, negative_prompts_2):
     if hasattr(p, 'keep_prompts'):
         return prompts, negative_prompts, prompts_2, negative_prompts_2
+
     if type(prompts) is str:
         prompts = [prompts]
     if type(negative_prompts) is str:
         negative_prompts = [negative_prompts]
+
     if hasattr(p, '[init_images]') and p.init_images is not None and len(p.init_images) > 1:
         while len(prompts) < len(p.init_images):
             prompts.append(prompts[-1])
         while len(negative_prompts) < len(p.init_images):
             negative_prompts.append(negative_prompts[-1])
+
+    while len(prompts) < p.batch_size:
+        prompts.append(prompts[-1])
+    while len(negative_prompts) < p.batch_size:
+        negative_prompts.append(negative_prompts[-1])
+
     while len(negative_prompts) < len(prompts):
         negative_prompts.append(negative_prompts[-1])
     while len(prompts) < len(negative_prompts):
         prompts.append(prompts[-1])
+
     if type(prompts_2) is str:
         prompts_2 = [prompts_2]
     if type(prompts_2) is list:
@@ -457,9 +466,7 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start):
         cls = shared.sd_model.__class__.__name__
         if cls in sd_models.i2i_pipes:
             steps = p.steps
-        elif 'Flex' in cls:
-            steps = p.steps
-        elif 'HiDreamImageEditingPipeline' in cls:
+        elif 'Flex' in cls or 'HiDreamImageEditingPipeline' in cls or 'Kontext' in cls:
             steps = p.steps
         elif use_denoise_start and (shared.sd_model_type == 'sdxl'):
             steps = p.steps // (1 - p.refiner_start)
@@ -542,7 +549,8 @@ def set_latents(p):
 def apply_circular(enable: bool, model):
     if not hasattr(model, 'unet') or not hasattr(model, 'vae'):
         return
-    if getattr(model, 'texture_tiling', False) == enable:
+    current = getattr(model, 'texture_tiling', None)
+    if isinstance(current, bool) and current == enable:
         return
     try:
         i = 0
@@ -553,7 +561,8 @@ def apply_circular(enable: bool, model):
             i += 1
             layer.padding_mode = 'circular' if enable else 'zeros'
         model.texture_tiling = enable
-        shared.log.debug(f'Apply texture tiling: enabled={enable} layers={i} cls={model.__class__.__name__} ')
+        if current is not None or enable:
+            shared.log.debug(f'Apply texture tiling: enabled={enable} layers={i} cls={model.__class__.__name__} ')
     except Exception as e:
         debug(f"Diffusers tiling failed: {e}")
 
@@ -572,10 +581,7 @@ def update_sampler(p, sd_model, second_pass=False):
     if hasattr(sd_model, 'scheduler'):
         if sampler_selection == 'None':
             return
-        if sampler_selection is None:
-            sampler = sd_samplers.all_samplers_map.get("UniPC")
-        else:
-            sampler = sd_samplers.all_samplers_map.get(sampler_selection, None)
+        sampler = sd_samplers.find_sampler(sampler_selection)
         if sampler is None:
             shared.log.warning(f'Sampler: sampler="{sampler_selection}" not found')
             sampler = sd_samplers.all_samplers_map.get("UniPC")
diff --git a/modules/processing_vae.py b/modules/processing_vae.py
index 36d5d0dda..b1af54d69 100644
--- a/modules/processing_vae.py
+++ b/modules/processing_vae.py
@@ -130,7 +130,7 @@ def full_vae_decode(latents, model):
     # normalize latents
     latents_mean = model.vae.config.get("latents_mean", None)
     latents_std = model.vae.config.get("latents_std", None)
-    scaling_factor = model.vae.config.get("scaling_factor", None)
+    scaling_factor = model.vae.config.get("scaling_factor", 1.0)
     shift_factor = model.vae.config.get("shift_factor", None)
     if latents_mean and latents_std:
         latents_mean = (torch.tensor(latents_mean).view(1, -1, 1, 1).to(latents.device, latents.dtype))
@@ -298,9 +298,9 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he
         latent_num_frames = (frames - 1) // model.vae_temporal_compression_ratio + 1
         latents = model._unpack_latents(latents.unsqueeze(0), latent_num_frames, height // 32, width // 32, model.transformer_spatial_patch_size, model.transformer_temporal_patch_size) # pylint: disable=protected-access
         latents = model._denormalize_latents(latents, model.vae.latents_mean, model.vae.latents_std, model.vae.config.scaling_factor) # pylint: disable=protected-access
-    if hasattr(model, '_unpack_latents') and hasattr(model, "vae_scale_factor") and width is not None and height is not None: # FLUX
+    if hasattr(model, '_unpack_latents') and hasattr(model, "vae_scale_factor") and width is not None and height is not None and latents.ndim == 3: # FLUX
         latents = model._unpack_latents(latents, height, width, model.vae_scale_factor) # pylint: disable=protected-access
-    if len(latents.shape) == 3: # lost a batch dim in hires
+    if latents.ndim == 3: # lost a batch dim in hires
         latents = latents.unsqueeze(0)
     if latents.shape[-1] <= 4: # not a latent, likely an image
         decoded = latents.float().cpu().numpy()
diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py
index ad80e3af0..d63fa4c12 100644
--- a/modules/prompt_parser_diffusers.py
+++ b/modules/prompt_parser_diffusers.py
@@ -7,7 +7,7 @@ import torch
 from compel.embeddings_provider import BaseTextualInversionManager, EmbeddingsProvider
 from transformers import PreTrainedTokenizer
 from modules import shared, prompt_parser, devices, sd_models
-from modules.prompt_parser_xhinker import get_weighted_text_embeddings_sd15, get_weighted_text_embeddings_sdxl_2p, get_weighted_text_embeddings_sd3, get_weighted_text_embeddings_flux1
+from modules.prompt_parser_xhinker import get_weighted_text_embeddings_sd15, get_weighted_text_embeddings_sdxl_2p, get_weighted_text_embeddings_sd3, get_weighted_text_embeddings_flux1, get_weighted_text_embeddings_chroma
 
 debug_enabled = os.environ.get('SD_PROMPT_DEBUG', None)
 debug = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
@@ -27,6 +27,7 @@ def prompt_compatible(pipe = None):
         'DemoFusion' not in pipe.__class__.__name__ and
         'StableCascade' not in pipe.__class__.__name__ and
         'Flux' not in pipe.__class__.__name__ and
+        'Chroma' not in pipe.__class__.__name__ and
         'HiDreamImage' not in pipe.__class__.__name__
     ):
         shared.log.warning(f"Prompt parser not supported: {pipe.__class__.__name__}")
@@ -53,7 +54,8 @@ class PromptEmbedder:
         self.negative_prompts = negative_prompts
         self.batchsize = len(self.prompts)
         self.attention = last_attention
-        self.allsame = self.compare_prompts()  # collapses batched prompts to single prompt if possible
+        self.allsame = False # dont collapse prompts
+        # self.allsame = self.compare_prompts()  # collapses batched prompts to single prompt if possible
         self.steps = steps
         self.clip_skip = clip_skip
         # All embeds are nested lists, outer list batch length, inner schedule length
@@ -61,6 +63,8 @@ class PromptEmbedder:
         self.positive_pooleds = [[]] * self.batchsize
         self.negative_prompt_embeds = [[]] * self.batchsize
         self.negative_pooleds = [[]] * self.batchsize
+        self.prompt_attention_masks = [[]] * self.batchsize
+        self.negative_prompt_attention_masks = [[]] * self.batchsize
         self.positive_schedule = None
         self.negative_schedule = None
         self.scheduled_prompt = False
@@ -83,7 +87,7 @@ class PromptEmbedder:
         self.checkcache(p)
         debug(f"Prompt encode: time={(time.time() - t0):.3f}")
 
-    def checkcache(self, p):
+    def checkcache(self, p) -> bool:
         if shared.opts.sd_textencoder_cache_size == 0:
             return False
         if self.scheduled_prompt:
@@ -108,13 +112,17 @@ class PromptEmbedder:
             if not any(flatten(emb) for emb in [self.prompt_embeds,
                                                 self.negative_prompt_embeds,
                                                 self.positive_pooleds,
-                                                self.negative_pooleds]):
+                                                self.negative_pooleds,
+                                                self.prompt_attention_masks,
+                                                self.negative_prompt_attention_masks]):
                 return False
             else:
                 cache[key] = {'prompt_embeds': self.prompt_embeds,
                               'negative_prompt_embeds': self.negative_prompt_embeds,
                               'positive_pooleds': self.positive_pooleds,
                               'negative_pooleds': self.negative_pooleds,
+                              'prompt_attention_masks': self.prompt_attention_masks,
+                              'negative_prompt_attention_masks': self.negative_prompt_attention_masks,
                               }
                 debug(f"Prompt cache: add={key}")
                 while len(cache) > int(shared.opts.sd_textencoder_cache_size):
@@ -128,6 +136,8 @@ class PromptEmbedder:
                 self.positive_pooleds = [self.positive_pooleds[0]] * self.batchsize
                 self.negative_prompt_embeds = [self.negative_prompt_embeds[0]] * self.batchsize
                 self.negative_pooleds = [self.negative_pooleds[0]] * self.batchsize
+                self.prompt_attention_masks = [self.prompt_attention_masks[0]] * self.batchsize
+                self.negative_prompt_attention_masks = [self.negative_prompt_attention_masks[0]] * self.batchsize
             debug(f"Prompt cache: get={key}")
             return True
 
@@ -166,24 +176,45 @@ class PromptEmbedder:
             self.positive_pooleds[batchidx].append(self.positive_pooleds[batchidx][idx])
         if len(self.negative_pooleds[batchidx]) > 0:
             self.negative_pooleds[batchidx].append(self.negative_pooleds[batchidx][idx])
+        if len(self.prompt_attention_masks[batchidx]) > 0:
+            self.prompt_attention_masks[batchidx].append(self.prompt_attention_masks[batchidx][idx])
+        if len(self.negative_prompt_attention_masks[batchidx]) > 0:
+            self.negative_prompt_attention_masks[batchidx].append(self.negative_prompt_attention_masks[batchidx][idx])
 
     def encode(self, pipe, positive_prompt, negative_prompt, batchidx):
         global last_attention # pylint: disable=global-statement
         self.attention = shared.opts.prompt_attention
         last_attention = self.attention
         if self.attention == "xhinker":
-            prompt_embed, positive_pooled, negative_embed, negative_pooled = get_xhinker_text_embeddings(pipe, positive_prompt, negative_prompt, self.clip_skip)
+            (
+                prompt_embed,
+                positive_pooled,
+                prompt_attention_mask,
+                negative_embed,
+                negative_pooled,
+                negative_prompt_attention_mask
+            ) = get_xhinker_text_embeddings(pipe, positive_prompt, negative_prompt, self.clip_skip)
         else:
-            prompt_embed, positive_pooled, negative_embed, negative_pooled = get_weighted_text_embeddings(pipe, positive_prompt, negative_prompt, self.clip_skip)
+            (
+                prompt_embed,
+                positive_pooled,
+                prompt_attention_mask,
+                negative_embed,
+                negative_pooled,
+                negative_prompt_attention_mask
+            ) = get_weighted_text_embeddings(pipe, positive_prompt, negative_prompt, self.clip_skip)
         if prompt_embed is not None:
-            self.prompt_embeds[batchidx].append(prompt_embed)
+            self.prompt_embeds[batchidx] = [prompt_embed]
         if negative_embed is not None:
-            self.negative_prompt_embeds[batchidx].append(negative_embed)
+            self.negative_prompt_embeds[batchidx] = [negative_embed]
         if positive_pooled is not None:
-            self.positive_pooleds[batchidx].append(positive_pooled)
+            self.positive_pooleds[batchidx] = [positive_pooled]
         if negative_pooled is not None:
-            self.negative_pooleds[batchidx].append(negative_pooled)
-
+            self.negative_pooleds[batchidx] = [negative_pooled]
+        if prompt_attention_mask is not None:
+            self.prompt_attention_masks[batchidx] = [prompt_attention_mask]
+        if negative_prompt_attention_mask is not None:
+            self.negative_prompt_attention_masks[batchidx] = [negative_prompt_attention_mask]
         if debug_enabled:
             get_tokens(pipe, 'positive', positive_prompt)
             get_tokens(pipe, 'negative', negative_prompt)
@@ -508,7 +539,11 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
 
     if "Flux" in pipe.__class__.__name__: # clip is only used for the pooled embeds
         prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=prompt, prompt_2=prompt_2, device=device, num_images_per_prompt=1)
-        return prompt_embeds, pooled_prompt_embeds, None, None # no negative support
+        return prompt_embeds, pooled_prompt_embeds, None, None, None, None # no negative support
+
+    if "Chroma" in pipe.__class__.__name__: # does not use clip and has no pooled embeds
+        prompt_embeds, _, prompt_attention_mask, negative_prompt_embeds, _, negative_prompt_attention_mask = pipe.encode_prompt(prompt=prompt, negative_prompt=neg_prompt, device=device, num_images_per_prompt=1)
+        return prompt_embeds, None, prompt_attention_mask, negative_prompt_embeds, None, negative_prompt_attention_mask
 
     if "HiDreamImage" in pipe.__class__.__name__: # clip is only used for the pooled embeds
         prompt_embeds_t5, negative_prompt_embeds_t5, prompt_embeds_llama3, negative_prompt_embeds_llama3, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
@@ -518,7 +553,7 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
         )
         prompt_embeds = [prompt_embeds_t5, prompt_embeds_llama3]
         negative_prompt_embeds = [negative_prompt_embeds_t5, negative_prompt_embeds_llama3]
-        return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds
+        return prompt_embeds, pooled_prompt_embeds, None, negative_prompt_embeds, negative_pooled_prompt_embeds, None
 
     if prompt != prompt_2:
         ps = [get_prompts_with_weights(pipe, p) for p in [prompt, prompt_2]]
@@ -631,7 +666,7 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
             negative_prompt_embeds, (0, t5_negative_prompt_embed.shape[-1] - negative_prompt_embeds.shape[-1])
         ).to(device)
         negative_prompt_embeds = torch.cat([negative_prompt_embeds, t5_negative_prompt_embed], dim=-2)
-    return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds
+    return prompt_embeds, pooled_prompt_embeds, None, negative_prompt_embeds, negative_pooled_prompt_embeds, None
 
 
 def get_xhinker_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int = None):
@@ -645,7 +680,7 @@ def get_xhinker_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", cl
         neg_prompt_2 = pipe.maybe_convert_prompt(neg_prompt_2, pipe.tokenizer_2)
     except Exception:
         pass
-    prompt_embed = positive_pooled = negative_embed = negative_pooled = None
+    prompt_embed = positive_pooled = negative_embed = negative_pooled = prompt_attention_mask = negative_prompt_attention_mask = None
 
     te1_device, te2_device, te3_device = None, None, None
     if hasattr(pipe, "text_encoder") and pipe.text_encoder.device != devices.device:
@@ -662,6 +697,8 @@ def get_xhinker_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", cl
         prompt_embed, negative_embed, positive_pooled, negative_pooled = get_weighted_text_embeddings_sd3(pipe=pipe, prompt=prompt, neg_prompt=neg_prompt, use_t5_encoder=bool(pipe.text_encoder_3))
     elif 'Flux' in pipe.__class__.__name__:
         prompt_embed, positive_pooled = get_weighted_text_embeddings_flux1(pipe=pipe, prompt=prompt, prompt2=prompt_2, device=devices.device)
+    elif 'Chroma' in pipe.__class__.__name__:
+        prompt_embed, prompt_attention_mask, negative_embed, negative_prompt_attention_mask = get_weighted_text_embeddings_chroma(pipe=pipe, prompt=prompt, neg_prompt=neg_prompt, device=devices.device)
     elif 'XL' in pipe.__class__.__name__:
         prompt_embed, negative_embed, positive_pooled, negative_pooled = get_weighted_text_embeddings_sdxl_2p(pipe=pipe, prompt=prompt, prompt_2=prompt_2, neg_prompt=neg_prompt, neg_prompt_2=neg_prompt_2)
     else:
@@ -674,4 +711,4 @@ def get_xhinker_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", cl
     if te3_device is not None:
         sd_models.move_model(pipe.text_encoder_3, te1_device, force=True)
 
-    return prompt_embed, positive_pooled, negative_embed, negative_pooled
+    return prompt_embed, positive_pooled, prompt_attention_mask, negative_embed, negative_pooled, negative_prompt_attention_mask
diff --git a/modules/prompt_parser_xhinker.py b/modules/prompt_parser_xhinker.py
index c0ddc9bc7..e2f40fc26 100644
--- a/modules/prompt_parser_xhinker.py
+++ b/modules/prompt_parser_xhinker.py
@@ -17,11 +17,13 @@
 ## -----------------------------------------------------------------------------
 
 import torch
+import torch.nn.functional as F
 from transformers import CLIPTokenizer, T5Tokenizer
 from diffusers import StableDiffusionPipeline
 from diffusers import StableDiffusionXLPipeline
 from diffusers import StableDiffusion3Pipeline
 from diffusers import FluxPipeline
+from diffusers import ChromaPipeline
 from modules.prompt_parser import parse_prompt_attention  # use built-in A1111 parser
 
 
@@ -86,8 +88,9 @@ def get_prompts_tokens_with_weights(
 
 
 def get_prompts_tokens_with_weights_t5(
-        t5_tokenizer: T5Tokenizer
-        , prompt: str
+        t5_tokenizer: T5Tokenizer,
+        prompt: str,
+        add_special_tokens: bool = True
 ):
     """
     Get prompt token ids and weights, this function works for both prompt and negative prompt
@@ -96,18 +99,22 @@ def get_prompts_tokens_with_weights_t5(
         prompt = "empty"
 
     texts_and_weights = parse_prompt_attention(prompt)
-    text_tokens, text_weights = [], []
+    text_tokens, text_weights, text_masks = [], [], []
     for word, weight in texts_and_weights:
         # tokenize and discard the starting and the ending token
-        token = t5_tokenizer(
-            word
-            , truncation=False  # so that tokenize whatever length prompt
-            , add_special_tokens=True
-        ).input_ids
-        # the returned token is a 1d list: [320, 1125, 539, 320]
+        inputs = t5_tokenizer(
+            word,
+            truncation=False,  # so that tokenize whatever length prompt
+            add_special_tokens=add_special_tokens,
+            return_length=False,
+        )
+
+        token = inputs.input_ids
+        mask = inputs.attention_mask
 
         # merge the new tokens to the all tokens holder: text_tokens
         text_tokens = [*text_tokens, *token]
+        text_masks = [*text_masks, *mask]
 
         # each token chunk will come with one weight, like ['red cat', 2.0]
         # need to expand weight for each token.
@@ -115,7 +122,7 @@ def get_prompts_tokens_with_weights_t5(
 
         # append the weight back to the weight holder: text_weights
         text_weights = [*text_weights, *chunk_weights]
-    return text_tokens, text_weights
+    return text_tokens, text_weights, text_masks
 
 
 def group_tokens_and_weights(
@@ -1068,11 +1075,11 @@ def get_weighted_text_embeddings_sd3(
     )
 
     # tokenizer 3
-    prompt_tokens_3, prompt_weights_3 = get_prompts_tokens_with_weights_t5(
+    prompt_tokens_3, prompt_weights_3, _ = get_prompts_tokens_with_weights_t5(
         pipe.tokenizer_3, prompt
     )
 
-    neg_prompt_tokens_3, neg_prompt_weights_3 = get_prompts_tokens_with_weights_t5(
+    neg_prompt_tokens_3, neg_prompt_weights_3, _ = get_prompts_tokens_with_weights_t5(
         pipe.tokenizer_3, neg_prompt
     )
 
@@ -1364,7 +1371,7 @@ def get_weighted_text_embeddings_flux1(
     )
 
     # tokenizer 2 - google/t5-v1_1-xxl
-    prompt_tokens_2, prompt_weights_2 = get_prompts_tokens_with_weights_t5(
+    prompt_tokens_2, prompt_weights_2, _ = get_prompts_tokens_with_weights_t5(
         pipe.tokenizer_2, prompt2
     )
 
@@ -1424,3 +1431,184 @@ def get_weighted_text_embeddings_flux1(
     t5_prompt_embeds = t5_prompt_embeds.to(dtype=pipe.text_encoder_2.dtype, device=device)
 
     return t5_prompt_embeds, prompt_embeds
+
+
+def get_weighted_text_embeddings_chroma(
+    pipe: ChromaPipeline,
+    prompt: str = "",
+    neg_prompt: str = "",
+    device=None
+):
+    """
+    This function can process long prompt with weights for Chroma model
+
+    Args:
+        pipe (ChromaPipeline)
+        prompt (str)
+        neg_prompt (str)
+        device (torch.device, optional): Device to run the embeddings on.
+    Returns:
+        prompt_embeds (torch.Tensor)
+        prompt_attention_mask (torch.Tensor)
+        neg_prompt_embeds (torch.Tensor)
+        neg_prompt_attention_mask (torch.Tensor)
+    """
+    if device is None:
+        device = pipe.text_encoder.device
+
+    dtype = pipe.text_encoder.dtype
+
+    prompt_tokens, prompt_weights, prompt_masks = get_prompts_tokens_with_weights_t5(
+        pipe.tokenizer, prompt, add_special_tokens=False
+    )
+
+    neg_prompt_tokens, neg_prompt_weights, neg_prompt_masks = get_prompts_tokens_with_weights_t5(
+        pipe.tokenizer, neg_prompt, add_special_tokens=False
+    )
+
+    padded_tokens, padded_weights, padded_masks = pad_prompt_tokens_to_same_size_chroma(
+        pipe,
+        [prompt_tokens, neg_prompt_tokens],
+        [prompt_weights, neg_prompt_weights],
+        [prompt_masks, neg_prompt_masks],
+        add_eos_token=True
+    )
+
+    prompt_tokens = padded_tokens[0]
+    prompt_weights = padded_weights[0]
+    prompt_masks = padded_masks[0]
+
+    prompt_embeds, prompt_masks = get_weighted_prompt_embeds_with_attention_mask_chroma(
+        pipe,
+        prompt_tokens,
+        prompt_weights,
+        prompt_masks,
+        device=device,
+        dtype=dtype)
+
+    neg_prompt_tokens = padded_tokens[1]
+    neg_prompt_weights = padded_weights[1]
+    neg_prompt_masks = padded_masks[1]
+
+    neg_prompt_embeds, neg_prompt_masks = get_weighted_prompt_embeds_with_attention_mask_chroma(
+        pipe,
+        neg_prompt_tokens,
+        neg_prompt_weights,
+        neg_prompt_masks,
+        device=device,
+        dtype=dtype)
+
+    # debug, will be removed later
+    prompt_with_mask, prompt_without_mask = debug_masked_tokens(pipe, prompt_tokens, prompt_masks.detach().tolist()[0])
+    neg_prompt_with_mask, neg_prompt_without_mask = debug_masked_tokens(pipe, neg_prompt_tokens, neg_prompt_masks.detach().tolist()[0])
+
+    return prompt_embeds, prompt_masks, neg_prompt_embeds, neg_prompt_masks
+
+
+# debug, will be removed later
+def debug_masked_tokens(pipe, prompt_tokens, prompt_masks):
+    prompt_with_mask = pipe.tokenizer.decode([token for token, mask in zip(prompt_tokens, prompt_masks) if mask == 1], skip_special_tokens=False)
+    prompt_without_mask = pipe.tokenizer.decode(prompt_tokens, skip_special_tokens=False)
+
+    return prompt_with_mask, prompt_without_mask
+
+
+def get_weighted_prompt_embeds_with_attention_mask_chroma(
+    pipe: ChromaPipeline,
+    tokens,
+    weights,
+    masks,
+    device,
+    dtype
+):
+    prompt_tokens = torch.tensor([tokens], dtype=torch.long, device=device)
+    prompt_masks = torch.tensor([masks], dtype=torch.long, device=device)
+    prompt_embeds = pipe.text_encoder(prompt_tokens, output_hidden_states=False, attention_mask=prompt_masks)[0].squeeze(0)
+    for z in range(len(weights)):
+        if weights[z] != 1.0:
+            prompt_embeds[z] = prompt_embeds[z] * weights[z]
+    prompt_embeds = prompt_embeds.unsqueeze(0).to(dtype=dtype, device=device)
+    return prompt_embeds, prompt_masks
+
+
+def pad_prompt_tokens_to_same_size_chroma(pipe, input_tokens, input_weights, input_masks, min_length=3, add_eos_token=True):
+    """
+    Implementation of Chroma's padding for prompt embeddings.
+    Pads the embeddings to the maximum length found in the batch, while ensuring
+    that the padding tokens are masked correctly while keeping at least one padding and one eos token unmasked.
+
+    https://huggingface.co/lodestones/Chroma#tldr-masking-t5-padding-tokens-enhanced-fidelity-and-increased-stability-during-training
+    """
+
+    input_tokens = input_tokens.copy()
+    input_weights = input_weights.copy()
+    input_masks = input_masks.copy()
+
+    pad_token_id = pipe.tokenizer.pad_token_id
+    eos_token_id = pipe.tokenizer.eos_token_id
+
+    for tokens, mask in zip(input_tokens, input_masks):
+        for j, token in enumerate(tokens):
+            if token == pad_token_id:
+                mask[j] = 0
+
+    max_token_count = max([len(x) for x in input_tokens] + [min_length])
+
+    padded_tokens = []
+    padded_weights = []
+    padded_masks = []
+
+    for tokens, weights, mask in zip(input_tokens, input_weights, input_masks):
+        current_length = len(tokens)
+
+        pad_length = 0
+
+        if current_length < max_token_count:
+            pad_length = max_token_count - current_length
+
+        elif pad_token_id not in tokens:
+            pad_length = 1
+
+        if pad_length > 0:
+            token_pad = [pad_token_id] * pad_length
+            weight_pad = [1.0] * pad_length
+            mask_pad = [0] * pad_length
+
+            tokens = tokens + token_pad
+            weights = weights + weight_pad
+            mask = mask + mask_pad
+
+        padded_tokens.append(tokens)
+        padded_weights.append(weights)
+        padded_masks.append(mask)
+
+    max_token_count = max([len(x) for x in padded_tokens])
+
+    for i, (tokens, weights, mask) in enumerate(zip(padded_tokens, padded_weights, padded_masks)):
+        if pad_token_id in tokens:
+            if tokens[-1] == pad_token_id:
+                mask[-1] = 1
+                continue
+
+        padded_tokens[i] = tokens + [pad_token_id]
+        padded_weights[i] = weights + [1.0]
+        padded_masks[i] = mask + [1]
+        max_token_count = max(max_token_count, len(padded_tokens[i]))
+
+    if add_eos_token:
+        max_token_count += 1  # eos token
+
+    for i in range(len(padded_tokens)):
+        if len(padded_tokens[i]) < max_token_count:
+            pad_length = max_token_count - len(padded_tokens[i])
+            padded_weights[i] += [1.0] * pad_length
+            padded_masks[i][-1] = 0
+            padded_masks[i] += [0] * (pad_length - 1) + [1]
+
+            if add_eos_token:
+                padded_tokens[i] += [pad_token_id] * (pad_length - 1) + [eos_token_id]
+                padded_masks[i][-2] = 1
+            else:
+                padded_tokens[i] += [pad_token_id] * pad_length
+
+    return padded_tokens, padded_weights, padded_masks
diff --git a/modules/sd_detect.py b/modules/sd_detect.py
index 3c2c1be72..be57edb4b 100644
--- a/modules/sd_detect.py
+++ b/modules/sd_detect.py
@@ -70,7 +70,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
             if 'lumina-next' in f.lower():
                 guess = 'Lumina-Next'
             if 'lumina-image-2' in f.lower():
-                guess = 'Lumina2'
+                guess = 'Lumina 2'
             if 'kolors' in f.lower():
                 guess = 'Kolors'
             if 'auraflow' in f.lower():
@@ -88,16 +88,23 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
             if 'omnigen' in f.lower():
                 guess = 'OmniGen'
                 pipeline = 'custom'
+            if 'omnigen2' in f.lower():
+                guess = 'OmniGen2'
+                pipeline = 'custom'
             if 'sd3' in f.lower():
                 guess = 'Stable Diffusion 3'
             if 'hidream' in f.lower():
                 guess = 'HiDream'
-            if 'flux' in f.lower() or 'flex.1' in f.lower() or 'lodestones' in f.lower():
+            if 'chroma' in f.lower():
+                guess = 'Chroma'
+            if 'flux' in f.lower() or 'flex.1' in f.lower():
                 guess = 'FLUX'
                 if size > 11000 and size < 16000:
                     warn(f'Model detected as FLUX UNET model, but attempting to load a base model: {op}={f} size={size} MB')
             if 'flex.2' in f.lower():
                 guess = 'FLEX'
+            if 'cosmos-predict2' in f.lower():
+                guess = 'Cosmos'
             # guess for diffusers
             index = os.path.join(f, 'model_index.json')
             if os.path.exists(index) and os.path.isfile(index):
@@ -113,6 +120,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                     guess = 'Stable Diffusion 3'
                 if callable(pipeline) and 'Lumina2' in pipeline.__name__:
                     guess = 'Lumina 2'
+
             # switch for specific variant
             if guess == 'Stable Diffusion' and 'inpaint' in f.lower():
                 guess = 'Stable Diffusion Inpaint'
@@ -122,6 +130,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                 guess = 'Stable Diffusion XL Inpaint'
             elif guess == 'Stable Diffusion XL' and 'instruct' in f.lower():
                 guess = 'Stable Diffusion XL Instruct'
+
             # get actual pipeline
             pipeline = shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline
             if debug_load is not None:
diff --git a/modules/sd_hijack_accelerate.py b/modules/sd_hijack_accelerate.py
index f8cf8983f..7f312a029 100644
--- a/modules/sd_hijack_accelerate.py
+++ b/modules/sd_hijack_accelerate.py
@@ -36,7 +36,7 @@ def hijack_set_module_tensor(
         # note: majority of time is spent on .to(old_value.dtype)
         if tensor_name in module._buffers: # pylint: disable=protected-access
             module._buffers[tensor_name] = value.to(device, old_value.dtype)  # pylint: disable=protected-access
-        elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device):  # pylint: disable=protected-access
+        elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device):  # pylint: disable=protected-access
             param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
             module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, old_value.dtype) # pylint: disable=protected-access
     t1 = time.time()
@@ -64,7 +64,7 @@ def hijack_set_module_tensor_simple(
     with devices.inference_context():
         if tensor_name in module._buffers: # pylint: disable=protected-access
             module._buffers[tensor_name] = value.to(device)  # pylint: disable=protected-access
-        elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device):  # pylint: disable=protected-access
+        elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device):  # pylint: disable=protected-access
             param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
             module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device) # pylint: disable=protected-access
     t1 = time.time()
diff --git a/modules/sd_hijack_hypertile.py b/modules/sd_hijack_hypertile.py
index c27bf8f39..ebe86948a 100644
--- a/modules/sd_hijack_hypertile.py
+++ b/modules/sd_hijack_hypertile.py
@@ -181,17 +181,19 @@ def context_hypertile_vae(p):
     if shared.opts.cross_attention_optimization == 'Sub-quadratic':
         shared.log.warning('Hypertile UNet is not compatible with Sub-quadratic cross-attention optimization')
         return nullcontext()
-    global height, width, max_h, max_w, error_reported # pylint: disable=global-statement
+    global max_h, max_w, error_reported # pylint: disable=global-statement
     error_reported = False
     error_reported = False
-    height, width = p.height, p.width
+    set_resolution(p)
     max_h, max_w = 0, 0
     vae = getattr(p.sd_model, "vae", None) if shared.native else getattr(p.sd_model, "first_stage_model", None)
+    if height == 0 or width == 0:
+        log.warning('Hypertile VAE disabled: resolution unknown')
+        return nullcontext()
     if height % 8 != 0 or width % 8 != 0:
         log.warning(f'Hypertile VAE disabled: width={width} height={height} are not divisible by 8')
         return nullcontext()
     if vae is None:
-        # shared.log.warning('Hypertile VAE is enabled but no VAE model was found')
         return nullcontext()
     else:
         tile_size = shared.opts.hypertile_vae_tile if shared.opts.hypertile_vae_tile > 0 else max(128, 64 * min(p.width // 128, p.height // 128))
@@ -208,11 +210,14 @@ def context_hypertile_unet(p):
     if shared.opts.cross_attention_optimization == 'Sub-quadratic' and not shared.cmd_opts.experimental:
         shared.log.warning('Hypertile UNet is not compatible with Sub-quadratic cross-attention optimization')
         return nullcontext()
-    global height, width, max_h, max_w, error_reported # pylint: disable=global-statement
+    global max_h, max_w, error_reported # pylint: disable=global-statement
     error_reported = False
-    height, width = p.height, p.width
+    set_resolution(p)
     max_h, max_w = 0, 0
     unet = getattr(p.sd_model, "unet", None) if shared.native else getattr(p.sd_model.model, "diffusion_model", None)
+    if height == 0 or width == 0:
+        log.warning('Hypertile VAE disabled: resolution unknown')
+        return nullcontext()
     if height % 8 != 0 or width % 8 != 0:
         log.warning(f'Hypertile UNet disabled: width={width} height={height} are not divisible by 8')
         return nullcontext()
@@ -229,17 +234,25 @@ def context_hypertile_unet(p):
 
 def hypertile_set(p, hr=False):
     from modules import shared
-    global height, width, error_reported, reset_needed, skip_hypertile # pylint: disable=global-statement
+    global error_reported, reset_needed, skip_hypertile # pylint: disable=global-statement
     if not shared.opts.hypertile_unet_enabled:
         return
     error_reported = False
+    set_resolution(p, hr=hr)
+    skip_hypertile = shared.opts.hypertile_hires_only and not getattr(p, 'is_hr_pass', False)
+    reset_needed = True
+
+
+def set_resolution(p, hr=False):
+    global height, width # pylint: disable=global-statement
     if hr:
         x = getattr(p, 'hr_upscale_to_x', 0)
         y = getattr(p, 'hr_upscale_to_y', 0)
         width = y if y > 0 else p.width
         height = x if x > 0 else p.height
     else:
-        width=p.width
-        height=p.height
-    skip_hypertile = shared.opts.hypertile_hires_only and not getattr(p, 'is_hr_pass', False)
-    reset_needed = True
+        width = p.width
+        height = p.height
+    if height == 0 or width == 0:
+        if hasattr(p, 'init_images') and isinstance(p.init_images, list) and len(p.init_images) > 0:
+            height, width = p.init_images[0].size
diff --git a/modules/sd_models.py b/modules/sd_models.py
index d8aad13a3..e998abefd 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -37,6 +37,7 @@ pipe_switch_task_exclude = [
     'InstantIRPipeline',
     'LTXConditionPipeline',
     'OmniGenPipeline',
+    'OmniGen2Pipeline',
     'PhotoMakerStableDiffusionXLPipeline',
     'PixelSmithXLPipeline',
     'StableDiffusion3ControlNetPipeline',
@@ -50,6 +51,7 @@ i2i_pipes = [
     'LEditsPPPipelineStableDiffusion',
     'LEditsPPPipelineStableDiffusionXL',
     'OmniGenPipeline',
+    'OmniGen2Pipeline',
     'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline',
     'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline',
 ]
@@ -156,6 +158,12 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True,
             shared.log.quiet(quiet, f'Setting {op}: fused-qkv=True')
         except Exception as e:
             shared.log.error(f'Setting {op}: fused-qkv=True {e}')
+    if shared.opts.enable_linfusion:
+        try:
+            from modules import linfusion
+            linfusion.apply(sd_model)
+        except Exception as e:
+            shared.log.error(f'Setting {op}: LinFusion=True {e}')
     if shared.opts.diffusers_eval:
         def eval_model(model, op=None, sd_model=None): # pylint: disable=unused-argument
             if hasattr(model, "requires_grad_"):
@@ -293,69 +301,99 @@ def load_diffuser_initial(diffusers_load_config, op='model'):
 
 def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='model'):
     sd_model = None
+    allow_post_quant = True
     unload_model_weights(op=op)
     shared.sd_model = None
     try:
         if model_type in ['Stable Cascade']: # forced pipeline
             from modules.model_stablecascade import load_cascade_combined
             sd_model = load_cascade_combined(checkpoint_info, diffusers_load_config)
+            allow_post_quant = True
         elif model_type in ['InstaFlow']: # forced pipeline
             pipeline = diffusers.utils.get_class_from_dynamic_module('instaflow_one_step', module_file='pipeline.py')
             shared_items.pipelines['InstaFlow'] = pipeline
             sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+            allow_post_quant = True
         elif model_type in ['SegMoE']: # forced pipeline
             from modules.segmoe.segmoe_model import SegMoEPipeline
             sd_model = SegMoEPipeline(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
             sd_model = sd_model.pipe # segmoe pipe does its stuff in __init__ and __call__ is the original pipeline
+            allow_post_quant = True
             shared_items.pipelines['SegMoE'] = SegMoEPipeline
         elif model_type in ['PixArt Sigma']: # forced pipeline
             from modules.model_pixart import load_pixart
             sd_model = load_pixart(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Sana']: # forced pipeline
             from modules.model_sana import load_sana
             sd_model = load_sana(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Lumina-Next']: # forced pipeline
             from modules.model_lumina import load_lumina
             sd_model = load_lumina(checkpoint_info, diffusers_load_config)
+            allow_post_quant = True
         elif model_type in ['Kolors']: # forced pipeline
             from modules.model_kolors import load_kolors
             sd_model = load_kolors(checkpoint_info, diffusers_load_config)
+            allow_post_quant = True
         elif model_type in ['AuraFlow']: # forced pipeline
             from modules.model_auraflow import load_auraflow
             sd_model = load_auraflow(checkpoint_info, diffusers_load_config)
+            allow_post_quant = True
         elif model_type in ['FLUX']:
             from modules.model_flux import load_flux
             sd_model = load_flux(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['FLEX']:
             from modules.model_flex import load_flex
             sd_model = load_flex(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
+        elif model_type in ['Chroma']:
+            from modules.model_chroma import load_chroma
+            sd_model = load_chroma(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Lumina 2']:
             from modules.model_lumina import load_lumina2
             sd_model = load_lumina2(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Stable Diffusion 3']:
             from modules.model_sd3 import load_sd3
             sd_model = load_sd3(checkpoint_info, cache_dir=shared.opts.diffusers_dir, config=diffusers_load_config.get('config', None))
+            allow_post_quant = False
         elif model_type in ['CogView 3']: # forced pipeline
             from modules.model_cogview import load_cogview3
             sd_model = load_cogview3(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['CogView 4']: # forced pipeline
             from modules.model_cogview import load_cogview4
             sd_model = load_cogview4(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['Meissonic']: # forced pipeline
             from modules.model_meissonic import load_meissonic
             sd_model = load_meissonic(checkpoint_info, diffusers_load_config)
+            allow_post_quant = True
+        elif model_type in ['OmniGen2']: # forced pipeline
+            from modules.model_omnigen2 import load_omnigen2
+            sd_model = load_omnigen2(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['OmniGen']: # forced pipeline
             from modules.model_omnigen import load_omnigen
             sd_model = load_omnigen(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
         elif model_type in ['HiDream']:
             from modules.model_hidream import load_hidream
             sd_model = load_hidream(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
+        elif model_type in ['Cosmos']:
+            from modules.model_cosmos import load_cosmos_t2i
+            sd_model = load_cosmos_t2i(checkpoint_info, diffusers_load_config)
+            allow_post_quant = False
     except Exception as e:
         shared.log.error(f'Load {op}: path="{checkpoint_info.path}" {e}')
         if debug_load:
             errors.display(e, 'Load')
-        return None
-    return sd_model
+        return None, True
+    return sd_model, allow_post_quant
 
 
 def load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_config, op='model'):
@@ -427,7 +465,6 @@ def load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_c
 
 def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_config, op='model'):
     sd_model = None
-    diffusers_load_config["local_files_only"] = diffusers_version < 28 # must be true for old diffusers, otherwise false but we override config for sd15/sdxl
     diffusers_load_config["extract_ema"] = shared.opts.diffusers_extract_ema
     if pipeline is None:
         shared.log.error(f'Load {op}: pipeline={shared.opts.diffusers_pipeline} not initialized')
@@ -536,6 +573,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             return
 
     sd_model = None
+    allow_post_quant = True
     try:
         # initial load only
         if sd_model is None:
@@ -567,7 +605,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
 
         # load with custom loader
         if sd_model is None:
-            sd_model = load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op)
+            sd_model, allow_post_quant = load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op)
             if sd_model is not None and not sd_model:
                 shared.log.error(f'Load {op}: type="{model_type}" pipeline="{pipeline}" not loaded')
                 return
@@ -620,7 +658,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         prompt_parser_diffusers.cache.clear()
 
         set_diffuser_options(sd_model, vae, op, offload=False)
-        sd_model = model_quant.do_post_load_quant(sd_model) # run this before move model so it can be compressed in CPU
+        sd_model = model_quant.do_post_load_quant(sd_model, allow=allow_post_quant) # run this before move model so it can be compressed in CPU
         timer.record("options")
 
         set_diffuser_offload(sd_model, op)
@@ -642,11 +680,6 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
             sd_model = sd_models_compile.compile_diffusers(sd_model)
         timer.record("compile")
 
-        if shared.opts.enable_linfusion:
-            from modules import linfusion
-            linfusion.apply(sd_model)
-            timer.record("linfusion")
-
     except Exception as e:
         shared.log.error(f"Load {op}: {e}")
         errors.display(e, "Model")
@@ -660,7 +693,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         from modules import modelstats
         modelstats.analyze()
 
-    shared.log.info(f"Load {op}: time={timer.summary()} native={get_native(sd_model)} memory={memory_stats()}")
+    shared.log.info(f"Load {op}: family={shared.sd_model_type} time={timer.dct()} native={get_native(sd_model)} memory={memory_stats()}")
 
 
 class DiffusersTaskType(Enum):
@@ -825,6 +858,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     sd_checkpoint_info = getattr(pipe, "sd_checkpoint_info", None)
     sd_model_checkpoint = getattr(pipe, "sd_model_checkpoint", None)
     embedding_db = getattr(pipe, "embedding_db", None)
+    loaded_loras = getattr(pipe, "loaded_loras", None)
     sd_model_hash = getattr(pipe, "sd_model_hash", None)
     has_accelerate = getattr(pipe, "has_accelerate", None)
     current_attn_name = getattr(pipe, "current_attn_name", None)
@@ -877,6 +911,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     new_pipe.has_accelerate = has_accelerate
     new_pipe.current_attn_name = current_attn_name
     new_pipe.default_scheduler = default_scheduler
+    new_pipe.loaded_loras = loaded_loras if loaded_loras is not None else {}
     if image_encoder is not None:
         new_pipe.image_encoder = image_encoder
     if feature_extractor is not None:
@@ -920,7 +955,7 @@ def set_diffusers_attention(pipe, quiet:bool=False):
     # if hasattr(pipe, 'pipe'):
     #    set_diffusers_attention(pipe.pipe)
 
-    if 'ControlNet' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")):
+    if 'Control' in pipe.__class__.__name__ or 'Adapter' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")):
         if shared.opts.cross_attention_optimization not in {"Scaled-Dot-Product", "Disabled"}:
             shared.log.warning(f"Attention: {shared.opts.cross_attention_optimization} is not compatible with {pipe.__class__.__name__}")
         else:
@@ -1086,7 +1121,8 @@ def clear_caches():
         lora_common.loaded_networks.clear()
         lora_common.previously_loaded_networks.clear()
         lora_load.lora_cache.clear()
-    from modules import prompt_parser_diffusers, memstats
+    from modules import prompt_parser_diffusers, memstats, sd_offload
+    sd_offload.offload_hook_instance = None
     prompt_parser_diffusers.cache.clear()
     memstats.reset_stats()
 
diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py
index 546297d25..9ef0024b7 100644
--- a/modules/sd_models_utils.py
+++ b/modules/sd_models_utils.py
@@ -151,14 +151,13 @@ def patch_diffuser_config(sd_model, model_file):
 
 
 def apply_function_to_model(sd_model, function, options, op=None):
-    if "Model" in options or "Transformer" in options:
-        if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'):
-            sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model)
     if "Model" in options:
         if hasattr(sd_model, 'model') and (hasattr(sd_model.model, 'config') or isinstance(sd_model.model, torch.nn.Module)):
             sd_model.model = function(sd_model.model, op="model", sd_model=sd_model)
         if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'):
             sd_model.unet = function(sd_model.unet, op="unet", sd_model=sd_model)
+        if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'):
+            sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model)
         if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model, 'decoder'):
             sd_model.decoder = None
             sd_model.decoder = sd_model.decoder_pipe.decoder = function(sd_model.decoder_pipe.decoder, op="decoder_pipe.decoder", sd_model=sd_model)
@@ -180,6 +179,8 @@ def apply_function_to_model(sd_model, function, options, op=None):
             sd_model.text_encoder_3 = function(sd_model.text_encoder_3, op="text_encoder_3", sd_model=sd_model)
         if hasattr(sd_model, 'text_encoder_4') and hasattr(sd_model.text_encoder_4, 'config'):
             sd_model.text_encoder_4 = function(sd_model.text_encoder_4, op="text_encoder_4", sd_model=sd_model)
+        if hasattr(sd_model, 'mllm') and hasattr(sd_model.mllm, 'config'):
+            sd_model.mllm = function(sd_model.mllm, op="text_encoder_mllm", sd_model=sd_model)
         if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder') and hasattr(sd_model.prior_pipe.text_encoder, 'config'):
             sd_model.prior_pipe.text_encoder = function(sd_model.prior_pipe.text_encoder, op="prior_pipe.text_encoder", sd_model=sd_model)
     if "VAE" in options:
diff --git a/modules/sd_offload.py b/modules/sd_offload.py
index c57165fb5..037f46ca2 100644
--- a/modules/sd_offload.py
+++ b/modules/sd_offload.py
@@ -4,6 +4,7 @@ import time
 import inspect
 import torch
 import accelerate.hooks
+import accelerate.utils.modeling
 from installer import log
 from modules import shared, devices, errors, model_quant
 from modules.timer import process as process_timer
@@ -11,10 +12,20 @@ from modules.timer import process as process_timer
 
 debug = os.environ.get('SD_MOVE_DEBUG', None) is not None
 debug_move = log.trace if debug else lambda *args, **kwargs: None
-offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4']
+offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'omnigen2', 'cogview4', 'cosmos', 'chroma']
 offload_post = ['h1']
 offload_hook_instance = None
-balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline']
+balanced_offload_exclude = ['CogView4Pipeline']
+accelerate_dtype_byte_size = None
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    try:
+        if dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz]:
+            dtype = accelerate.utils.modeling.CustomDtype.FP8
+    except Exception: # catch since older torch many not have defined dtypes
+        pass
+    return accelerate_dtype_byte_size(dtype)
 
 
 def get_signature(cls):
@@ -58,6 +69,7 @@ def set_accelerate(sd_model):
 
 
 def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
+    global accelerate_dtype_byte_size # pylint: disable=global-statement
     t0 = time.time()
     if not shared.native:
         shared.log.warning('Attempting to use offload with backend=original')
@@ -67,6 +79,9 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
         return
     if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate):
         sd_model.has_accelerate = False
+    if accelerate_dtype_byte_size is None:
+        accelerate_dtype_byte_size = accelerate.utils.modeling.dtype_byte_size
+        accelerate.utils.modeling.dtype_byte_size = dtype_byte_size
     if shared.opts.diffusers_offload_mode == "none":
         if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type:
             shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model')
@@ -156,21 +171,25 @@ class OffloadHook(accelerate.hooks.ModelHook):
         return module
 
     def pre_forward(self, module, *args, **kwargs):
-        if devices.normalize_device(module.device) != devices.normalize_device(devices.device):
+        if not devices.same_device(module.device, devices.device):
             device_index = torch.device(devices.device).index
             if device_index is None:
                 device_index = 0
             max_memory = { device_index: self.gpu, "cpu": self.cpu }
             device_map = getattr(module, "balanced_offload_device_map", None)
             if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None):
+                # try:
                 device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory)
+                # except Exception as e:
+                #     shared.log.error(f'Offload: type=balanced module={module.__class__.__name__} {e}')
             offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
             if devices.backend == "directml":
                 keys = device_map.keys()
                 for v in keys:
                     if isinstance(device_map[v], int):
                         device_map[v] = f"{devices.device.type}:{device_map[v]}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type
-            module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
+            if device_map is not None:
+                module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
             module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
             module.balanced_offload_device_map = device_map
             module.balanced_offload_max_memory = max_memory
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index f73520027..3be1d2537 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -12,10 +12,21 @@ samplers = all_samplers
 samplers_for_img2img = all_samplers
 samplers_map = {}
 loaded_config = None
-flow_models =  ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'CogView4', 'HiDream']
+flow_models = ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'CogView4', 'HiDream', 'Chroma']
 flow_models += ['Hunyuan', 'LTX', 'Mochi']
 
 
+def find_sampler(name:str):
+    if name is None or name == 'None':
+        return all_samplers_map.get("UniPC", None)
+    for sampler in all_samplers:
+        if sampler.name.lower() == name.lower() or name in sampler.aliases:
+            debug(f'Find sampler: name="{name}" found={sampler.name}')
+            return sampler
+    debug(f'Find sampler: name="{name}" found=None')
+    return None
+
+
 def list_samplers():
     global all_samplers # pylint: disable=global-statement
     global all_samplers_map # pylint: disable=global-statement
diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index 52fd6313e..09a06352a 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -9,7 +9,7 @@ from modules import shared, devices, processing, images, sd_vae_approx, sd_vae_t
 
 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
 approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE": 3 }
-flow_models = ['f1', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2', 'cogview4', 'h1']
+flow_models = ['f1', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2', 'cogview4', 'h1', 'cosmos', 'chroma', 'omnigen', 'omnigen2']
 warned = False
 queue_lock = threading.Lock()
 
@@ -52,9 +52,9 @@ def single_sample_to_image(sample, approximation=None):
         if len(sample.shape) == 4 and sample.shape[0]: # likely animatediff latent
             sample = sample.permute(1, 0, 2, 3)[0]
         if approximation == 2: # TAESD
-            if (len(sample.shape) == 3 or len(sample.shape) == 4) and shared.opts.live_preview_downscale and (sample.shape[-1] > 128 or sample.shape[-2] > 128):
+            if (len(sample.shape) == 3 or len(sample.shape) == 4) and shared.opts.live_preview_downscale and (sample.shape[-1]*sample.shape[-2] > 128*128):
                 try:
-                    scale = 128 / max(sample.shape[-1], sample.shape[-2])
+                    scale = (128 * 128) / (sample.shape[-1] * sample.shape[-2])
                     sample = torch.nn.functional.interpolate(sample.unsqueeze(0), scale_factor=[scale, scale], mode='bilinear', align_corners=False)[0]
                 except Exception:
                     pass
diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py
index f151ab151..8df2d0dbd 100644
--- a/modules/sd_samplers_diffusers.py
+++ b/modules/sd_samplers_diffusers.py
@@ -13,30 +13,30 @@ debug_log = shared.log.trace if debug else lambda *args, **kwargs: None
 try:
     from diffusers import (
         CMStochasticIterativeScheduler,
-        UniPCMultistepScheduler,
-        DDIMScheduler,
-        EulerDiscreteScheduler,
-        EulerAncestralDiscreteScheduler,
-        EDMEulerScheduler,
-        FlowMatchEulerDiscreteScheduler,
-        DEISMultistepScheduler,
-        SASolverScheduler,
-        DPMSolverSinglestepScheduler,
-        DPMSolverMultistepScheduler,
-        DPMSolverMultistepInverseScheduler,
-        EDMDPMSolverMultistepScheduler,
         CosineDPMSolverMultistepScheduler,
-        DPMSolverSDEScheduler,
-        HeunDiscreteScheduler,
-        FlowMatchHeunDiscreteScheduler,
-        LCMScheduler,
-        FlowMatchLCMScheduler,
-        PNDMScheduler,
-        IPNDMScheduler,
+        DDIMScheduler,
         DDPMScheduler,
-        LMSDiscreteScheduler,
-        KDPM2DiscreteScheduler,
+        DEISMultistepScheduler,
+        DPMSolverMultistepInverseScheduler,
+        DPMSolverMultistepScheduler,
+        DPMSolverSDEScheduler,
+        DPMSolverSinglestepScheduler,
+        EDMDPMSolverMultistepScheduler,
+        EDMEulerScheduler,
+        EulerAncestralDiscreteScheduler,
+        EulerDiscreteScheduler,
+        FlowMatchEulerDiscreteScheduler,
+        FlowMatchHeunDiscreteScheduler,
+        FlowMatchLCMScheduler,
+        HeunDiscreteScheduler,
+        IPNDMScheduler,
         KDPM2AncestralDiscreteScheduler,
+        KDPM2DiscreteScheduler,
+        LCMScheduler,
+        LMSDiscreteScheduler,
+        PNDMScheduler,
+        SASolverScheduler,
+        UniPCMultistepScheduler,
     )
 except Exception as e:
     shared.log.error(f'Sampler import: version={diffusers.__version__} error: {e}')
@@ -52,10 +52,6 @@ try:
     from modules.schedulers.scheduler_ufogen import UFOGenScheduler # pylint: disable=ungrouped-imports
     from modules.schedulers.scheduler_unipc_flowmatch import FlowUniPCMultistepScheduler # pylint: disable=ungrouped-imports
     from modules.perflow import PeRFlowScheduler # pylint: disable=ungrouped-imports
-    # from modules.schedulers.scheduler_kohaku import KohakuLoNyuYogScheduler # pylint: disable=ungrouped-imports
-    # from modules.schedulers.scheduler_smea import SMEAScheduler # pylint: disable=ungrouped-imports
-    # from modules.schedulers.scheduler_dy import DYScheduler # pylint: disable=ungrouped-imports
-    # from modules.schedulers.scheduler_negative import EulerNegativeScheduler # pylint: disable=ungrouped-imports
 except Exception as e:
     shared.log.error(f'Sampler import: version={diffusers.__version__} error: {e}')
     if os.environ.get('SD_SAMPLER_DEBUG', None) is not None:
@@ -74,10 +70,6 @@ config = {
     'Euler SGM': { 'steps_offset': 0, 'interpolation_type': "linear", 'rescale_betas_zero_snr': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'trailing', 'use_beta_sigmas': False, 'use_exponential_sigmas': False, 'use_karras_sigmas': False, 'prediction_type': "sample" },
     'Euler EDM': { 'sigma_schedule': "karras" },
     'Euler FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1, 'use_dynamic_shifting': False, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False },
-    # 'Euler SMEA': {},
-    # 'Euler DY': {},
-    # 'Euler Negative': {},
-    # 'Kohaku LoNyu': {},
 
     'DPM++': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_flow_sigmas': False, 'use_beta_sigmas': False, 'use_lu_lambdas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 1 },
     'DPM++ 2M': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_flow_sigmas': False, 'use_beta_sigmas': False, 'use_lu_lambdas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 2 },
@@ -135,9 +127,6 @@ samplers_data_diffusers = [
     SamplerData('Euler SGM', lambda model: DiffusionSampler('Euler SGM', EulerDiscreteScheduler, model), [], {}),
     SamplerData('Euler EDM', lambda model: DiffusionSampler('Euler EDM', EDMEulerScheduler, model), [], {}),
     SamplerData('Euler FlowMatch', lambda model: DiffusionSampler('Euler FlowMatch', FlowMatchEulerDiscreteScheduler, model), [], {}),
-    # SamplerData('Euler SMEA', lambda model: DiffusionSampler('Euler SMEA', SMEAScheduler, model), [], {}),
-    # SamplerData('Euler DY', lambda model: DiffusionSampler('Euler DY', DYScheduler, model), [], {}),
-    # SamplerData('Euler Negative', lambda model: DiffusionSampler('Euler Negative', EulerNegativeScheduler, model), [], {}),
 
     SamplerData('DPM++', lambda model: DiffusionSampler('DPM++', DPMSolverMultistepScheduler, model), [], {}),
     SamplerData('DPM++ 2M', lambda model: DiffusionSampler('DPM++ 2M', DPMSolverMultistepScheduler, model), [], {}),
@@ -185,7 +174,6 @@ samplers_data_diffusers = [
     SamplerData('TDD', lambda model: DiffusionSampler('TDD', TDDScheduler, model), [], {}),
     SamplerData('PeRFlow', lambda model: DiffusionSampler('PeRFlow', PeRFlowScheduler, model), [], {}),
     SamplerData('UFOGen', lambda model: DiffusionSampler('UFOGen', UFOGenScheduler, model), [], {}),
-    # SamplerData('Kohaku LoNyu', lambda model: DiffusionSampler('Kohaku LoNyu', KohakuLoNyuYogScheduler, model), [], {}),
 
     SamplerData('Same as primary', None, [], {}),
 ]
diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py
index 741d349bc..d0ccd6ac3 100644
--- a/modules/sd_vae_remote.py
+++ b/modules/sd_vae_remote.py
@@ -12,14 +12,27 @@ hf_decode_endpoints = {
     'sd': 'https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud',
     'sdxl': 'https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud',
     'f1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
-    'h1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
     'hunyuanvideo': 'https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud',
 }
+hf_decode_endpoints['pixartalpha'] = hf_decode_endpoints['sd']
+hf_decode_endpoints['pixartsigma'] = hf_decode_endpoints['sdxl']
+hf_decode_endpoints['omnigen'] = hf_decode_endpoints['sdxl']
+hf_decode_endpoints['h1'] = hf_decode_endpoints['f1']
+hf_decode_endpoints['chroma'] = hf_decode_endpoints['f1']
+hf_decode_endpoints['lumina2'] = hf_decode_endpoints['f1']
+
 hf_encode_endpoints = {
     'sd': 'https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud',
     'sdxl': 'https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud',
     'f1': 'https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud',
+    'chroma': 'https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud',
 }
+hf_encode_endpoints['pixartalpha'] = hf_encode_endpoints['sd']
+hf_encode_endpoints['pixartsigma'] = hf_encode_endpoints['sdxl']
+hf_encode_endpoints['omnigen'] = hf_encode_endpoints['sdxl']
+hf_encode_endpoints['h1'] = hf_encode_endpoints['f1']
+hf_encode_endpoints['lumina2'] = hf_encode_endpoints['f1']
+
 dtypes = {
     "float16": torch.float16,
     "float32": torch.float32,
@@ -48,7 +61,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
         params = {}
         try:
             latent = latent_copy[i]
-            if model_type != 'f1':
+            if model_type not in ['f1', 'chroma']:
                 latent = latent.unsqueeze(0)
             params = {
                 "input_tensor_type": "binary",
@@ -74,7 +87,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
                 params["output_type"] = "pt"
                 params["output_tensor_type"] = "binary"
                 headers["Accept"] = "tensor/binary"
-            if (model_type == 'f1' or model_type == 'h1') and (width > 0) and (height > 0):
+            if model_type in {'f1', 'h1', 'lumina2', 'chroma'} and (width > 0) and (height > 0):
                 params['width'] = width
                 params['height'] = height
             if shared.sd_model.vae is not None and shared.sd_model.vae.config is not None:
diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py
index 3bcd1b322..3ef7e123d 100644
--- a/modules/sd_vae_taesd.py
+++ b/modules/sd_vae_taesd.py
@@ -36,7 +36,7 @@ prev_cls = ''
 prev_type = ''
 prev_model = ''
 lock = threading.Lock()
-supported = ['sd', 'sdxl', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanvideo', 'mochivideo', 'pixartsigma', 'pixartalpha']
+supported = ['sd', 'sdxl', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanvideo', 'mochivideo', 'pixartsigma', 'pixartalpha', 'omnigen']
 
 
 def warn_once(msg, variant=None):
@@ -52,34 +52,36 @@ def warn_once(msg, variant=None):
 def get_model(model_type = 'decoder', variant = None):
     global prev_cls, prev_type, prev_model # pylint: disable=global-statement
     from modules import shared
-    cls = shared.sd_model_type
-    if cls in {'ldm', 'pixartalpha'}:
-        cls = 'sd'
-    elif cls in {'h1', 'lumina2'}:
-        cls = 'f1'
-    elif cls == 'pixartsigma':
-        cls = 'sdxl'
-    elif cls not in supported:
-        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+    model_cls = shared.sd_model_type
+    if model_cls is None or model_cls == 'none':
+        return None
+    elif model_cls in {'ldm', 'pixartalpha'}:
+        model_cls = 'sd'
+    elif model_cls in {'h1', 'lumina2', 'chroma'}:
+        model_cls = 'f1'
+    elif model_cls in {'pixartsigma', 'omnigen'}:
+        model_cls = 'sdxl'
+    elif model_cls not in supported:
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
     variant = variant or shared.opts.taesd_variant
     folder = os.path.join(paths.models_path, "TAESD")
     os.makedirs(folder, exist_ok=True)
     if variant.startswith('TAE'):
         cfg = TAESD_MODELS[variant]
-        if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
             return cfg['model']
-        fn = os.path.join(folder, cfg['fn'] + cls + '_' + model_type + '.pth')
+        fn = os.path.join(folder, cfg['fn'] + model_type + '_' + model_cls + '.pth')
         if not os.path.exists(fn):
             uri = cfg['uri']
             if not uri.endswith('.pth'):
-                uri += '/tae' + cls + '_' + model_type + '.pth'
+                uri += '/tae' + model_cls + '_' + model_type + '.pth'
             try:
                 shared.log.info(f'Decode: type="taesd" variant="{variant}": uri="{uri}" fn="{fn}" download')
                 torch.hub.download_url_to_file(uri, fn)
             except Exception as e:
                 warn_once(f'download uri={uri} {e}', variant=variant)
         if os.path.exists(fn):
-            prev_cls = cls
+            prev_cls = model_cls
             prev_type = model_type
             prev_model = variant
             shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" load')
@@ -97,14 +99,14 @@ def get_model(model_type = 'decoder', variant = None):
                 TAESD_MODELS[variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None)
             return TAESD_MODELS[variant]['model']
     elif variant.startswith('Hybrid'):
-        cfg = CQYAN_MODELS[variant].get(cls, None)
-        if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        cfg = CQYAN_MODELS[variant].get(model_cls, None)
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
             return cfg['model']
         if cfg is None:
-            warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+            warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
             return None
         repo = cfg['repo']
-        prev_cls = cls
+        prev_cls = model_cls
         prev_type = model_type
         prev_model = variant
         shared.log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load')
@@ -116,10 +118,12 @@ def get_model(model_type = 'decoder', variant = None):
             from modules.taesd.hybrid_small import AutoencoderSmall
             vae = AutoencoderSmall.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype)
         vae = vae.to(devices.device, dtype=dtype)
-        CQYAN_MODELS[variant][cls]['model'] = vae
+        CQYAN_MODELS[variant][model_cls]['model'] = vae
         return vae
+    elif variant is None:
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} variant is none', variant=variant)
     else:
-        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
     return None
 
 
diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py
index f5b7179a7..b4e251601 100644
--- a/modules/sdnq/__init__.py
+++ b/modules/sdnq/__init__.py
@@ -21,6 +21,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
         is_conv_transpose_type = False
         is_linear_type = False
         result_shape = None
+        original_shape = layer.weight.shape
         if torch_dtype is None:
             torch_dtype = devices.dtype
 
@@ -56,8 +57,11 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
             output_channel_size, channel_size = layer.weight.shape
             if use_quantized_matmul:
                 use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and channel_size >= 32 and output_channel_size >= 32
-                if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
-                    use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
+                if use_quantized_matmul:
+                    if dtype_dict[weights_dtype]["is_integer"]:
+                        use_quantized_matmul = output_channel_size % 8 == 0 and channel_size % 8 == 0
+                    else:
+                        use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
 
         if group_size == 0:
             if is_linear_type:
@@ -73,13 +77,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
                 num_of_groups = 1
             else:
                 num_of_groups = channel_size // group_size
-                while channel_size % group_size != 0: # find something divisible
+                while num_of_groups * group_size != channel_size: # find something divisible
                     num_of_groups -= 1
                     if num_of_groups <= 1:
                         group_size = channel_size
                         num_of_groups = 1
                         break
-                    group_size = channel_size / num_of_groups
+                    group_size = channel_size // num_of_groups
             group_size = int(group_size)
             num_of_groups = int(num_of_groups)
 
@@ -136,6 +140,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
             quantized_weight_shape=layer.weight.shape,
             result_dtype=torch_dtype,
             result_shape=result_shape,
+            original_shape=original_shape,
             weights_dtype=weights_dtype,
             use_quantized_matmul=use_quantized_matmul,
         )
@@ -144,15 +149,17 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
 
         layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul)
         layer.forward = layer.forward.__get__(layer, layer.__class__)
-        devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}")
+        #devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}")
     return layer
 
 
-def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument
+def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None, modules_to_not_convert: List[str] = []): # pylint: disable=unused-argument
     has_children = list(model.children())
     if not has_children:
         return model
     for module_param_name, module in model.named_children():
+        if module_param_name in modules_to_not_convert:
+            continue
         if hasattr(module, "weight") and module.weight is not None:
             module = sdnq_quantize_layer(
                 module,
@@ -179,6 +186,7 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
                 quantization_device=quantization_device,
                 return_device=return_device,
                 param_name=module_param_name,
+                modules_to_not_convert=modules_to_not_convert,
             )
     return model
 
@@ -200,7 +208,7 @@ def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, Li
     return scale
 
 
-def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str):
+def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
     if dtype_dict[weights_dtype]["is_unsigned"]:
         scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype)
         quantized_weight = torch.sub(weight, zero_point).div_(scale)
@@ -229,7 +237,7 @@ class SDNQQuantizer(DiffusersQuantizer):
     required_packages = None
     torch_dtype = None
 
-    def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation
+    def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
         self.modules_to_not_convert = []
 
@@ -381,7 +389,21 @@ class SDNQConfig(QuantizationConfigMixin):
         weights_dtype (`str`, *optional*, defaults to `"int8"`):
             The target dtype for the weights after quantization. Supported values are:
             ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
-       modules_to_not_convert (`list`, *optional*, default to `None`):
+        weights_dtype (`int`, *optional*, defaults to `0`):
+            Used to decide how many elements of a tensor will share the same quantization group.
+        quant_conv (`bool`, *optional*, defaults to `False`):
+            Enabling this option will quantize the convolutional layers in UNet models too.
+        use_quantized_matmul (`bool`, *optional*, defaults to `False`):
+            Enabling this option will use quantized INT8 or FP8 MatMul instead of BF16 / FP16.
+        use_quantized_matmul_conv (`bool`, *optional*, defaults to `False`):
+            Same as use_quantized_matmul_conv but for the convolutional layers with UNets like SDXL.
+        dequantize_fp32 (`bool`, *optional*, defaults to `False`):
+            Enabling this option will use FP32 on the dequantization step.
+        quantization_device (`torch.device`, *optional*, defaults to `None`):
+            Used to set which device will be used for the quantization calculation on model load.
+        return_device (`torch.device`, *optional*, defaults to `None`):
+            Used to set which device will the quantized weights be sent back to.
+        modules_to_not_convert (`list`, *optional*, default to `None`):
             The list of modules to not quantize, useful for quantizing models that explicitly require to have some
             modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
     """
@@ -419,25 +441,8 @@ class SDNQConfig(QuantizationConfigMixin):
         accepted_weights = ["int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"]
         if self.weights_dtype not in accepted_weights:
             raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")
-        if not isinstance(self.modules_to_not_convert, list):
+
+        if self.modules_to_not_convert is None:
+            self.modules_to_not_convert = []
+        elif not isinstance(self.modules_to_not_convert, list):
             self.modules_to_not_convert = [self.modules_to_not_convert]
-
-
-class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class
-    def __init__(self, T5DenseGatedActDense, dtype):
-        super().__init__()
-        self.wi_0 = T5DenseGatedActDense.wi_0
-        self.wi_1 = T5DenseGatedActDense.wi_1
-        self.wo = T5DenseGatedActDense.wo
-        self.dropout = T5DenseGatedActDense.dropout
-        self.act = T5DenseGatedActDense.act
-        self.torch_dtype = dtype
-
-    def forward(self, hidden_states):
-        hidden_gelu = self.act(self.wi_0(hidden_states))
-        hidden_linear = self.wi_1(hidden_states)
-        hidden_states = hidden_gelu * hidden_linear
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced
-        hidden_states = self.wo(hidden_states)
-        return hidden_states
diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py
index ccfdf1227..ab8b70c0d 100644
--- a/modules/sdnq/common.py
+++ b/modules/sdnq/common.py
@@ -2,33 +2,34 @@
 
 import sys
 import torch
-from accelerate.utils import CustomDtype
 from modules import devices
 
 torch_version = float(torch.__version__[:3])
 
 dtype_dict = {
     "int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True},
-    "int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
-    "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": "int7", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": "int6", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": "int5", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": "int4", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": "int3", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": "int2", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
     "uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "uint3": {"min": 0, "max": 7, "num_bits": 3, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
-    "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": "uint6", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": "uint5", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": "uint4", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint3": {"min": 0, "max": 7, "num_bits": 3, "target_dtype": "uint3", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": "uint2", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
     "uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True},
     "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False},
     "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False},
-    "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False},
-    "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False},
 }
 dtype_dict["bool"] = dtype_dict["uint1"]
+if hasattr(torch, "float8_e4m3fnuz"):
+    dtype_dict["float8_e4m3fnuz"] = {"min": -240, "max": 240, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False}
+if hasattr(torch, "float8_e5m2fnuz"):
+    dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False}
 
 use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
 quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2")
diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py
index be59019a3..45cedfa79 100644
--- a/modules/sdnq/dequantizer.py
+++ b/modules/sdnq/dequantizer.py
@@ -7,50 +7,52 @@ from .common import dtype_dict
 from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict
 
 
-def dequantize_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
-    result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype)
+def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
+    result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale).to(dtype=dtype)
     if result_shape is not None:
         result = result.reshape(result_shape)
     return result
 
 
-def dequantize_symmetric(input: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
+def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
     if skip_quantized_matmul:
-        result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
+        result = weight.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
     else:
-        result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
+        result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
     if result_shape is not None:
         result = result.reshape(result_shape)
     return result
 
 
-def dequantize_symmetric_with_bias(input: torch.CharTensor, bias: torch.FloatTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
-    return torch.addcmul(bias, input.to(dtype=scale.dtype), scale).to(dtype=dtype).reshape(result_shape)
+def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatTensor, bias: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
+    return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).reshape(result_shape)
 
 
-def dequantize_packed_int_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
-    return dequantize_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape)
+def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
+    return dequantize_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](weight, shape), scale, zero_point, dtype, result_shape)
 
 
-def dequantize_packed_int_symmetric(input: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
+def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
     if skip_quantized_matmul:
-        return dequantize_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
+        return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
     else:
-        return dequantize_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
+        return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
 
 
 class AsymmetricWeightsDequantizer(torch.nn.Module):
     def __init__(
         self,
-        scale: torch.Tensor,
-        zero_point: torch.Tensor,
+        scale: torch.FloatTensor,
+        zero_point: torch.FloatTensor,
         result_dtype: torch.dtype,
         result_shape: torch.Size,
+        original_shape: torch.Size,
         weights_dtype: str,
         **kwargs, # pylint: disable=unused-argument
     ):
         super().__init__()
         self.weights_dtype = weights_dtype
+        self.original_shape = original_shape
         self.use_quantized_matmul = False
         self.result_dtype = result_dtype
         self.result_shape = result_shape
@@ -67,15 +69,17 @@ class AsymmetricWeightsDequantizer(torch.nn.Module):
 class SymmetricWeightsDequantizer(torch.nn.Module):
     def __init__(
         self,
-        scale: torch.Tensor,
+        scale: torch.FloatTensor,
         result_dtype: torch.dtype,
         result_shape: torch.Size,
+        original_shape: torch.Size,
         weights_dtype: str,
         use_quantized_matmul: bool = False,
         **kwargs, # pylint: disable=unused-argument
     ):
         super().__init__()
         self.weights_dtype = weights_dtype
+        self.original_shape = original_shape
         self.use_quantized_matmul = use_quantized_matmul
         self.result_dtype = result_dtype
         self.result_shape = result_shape
@@ -91,17 +95,19 @@ class SymmetricWeightsDequantizer(torch.nn.Module):
 class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
     def __init__(
         self,
-        scale: torch.Tensor,
-        zero_point: torch.Tensor,
+        scale: torch.FloatTensor,
+        zero_point: torch.FloatTensor,
         quantized_weight_shape: torch.Size,
         result_dtype: torch.dtype,
         result_shape: torch.Size,
+        original_shape: torch.Size,
         weights_dtype: str,
         **kwargs, # pylint: disable=unused-argument
     ):
         super().__init__()
         self.weights_dtype = weights_dtype
         self.use_quantized_matmul = False
+        self.original_shape = original_shape
         self.quantized_weight_shape = quantized_weight_shape
         self.result_dtype = result_dtype
         self.result_shape = result_shape
@@ -118,16 +124,18 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
 class PackedINTSymmetricWeightsDequantizer(torch.nn.Module):
     def __init__(
         self,
-        scale: torch.Tensor,
+        scale: torch.FloatTensor,
         quantized_weight_shape: torch.Size,
         result_dtype: torch.dtype,
         result_shape: torch.Size,
+        original_shape: torch.Size,
         weights_dtype: str,
         use_quantized_matmul: bool = False,
         **kwargs, # pylint: disable=unused-argument
     ):
         super().__init__()
         self.weights_dtype = weights_dtype
+        self.original_shape = original_shape
         self.use_quantized_matmul = use_quantized_matmul
         self.quantized_weight_shape = quantized_weight_shape
         self.result_dtype = result_dtype
diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py
index 9caa12f7d..048354cca 100644
--- a/modules/sdnq/forward.py
+++ b/modules/sdnq/forward.py
@@ -95,7 +95,7 @@ def fp8_matmul_tensorwise(
     dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
     input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
     if bias is not None:
-        return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), bias, scale, return_dtype, output_shape)
+        return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, output_shape)
     else:
         return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape)
 
@@ -115,7 +115,7 @@ def int8_matmul(
     output_shape[-1] = weight.shape[-1]
     input, scale = quantize_int8_matmul_input(input, scale)
     if bias is not None:
-        return dequantize_symmetric_with_bias(torch._int_mm(input, weight), bias, scale, return_dtype, output_shape)
+        return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape)
     else:
         return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
 
@@ -236,7 +236,7 @@ def conv_fp8_matmul_tensorwise(
             result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype))
         result = torch.cat(result, dim=-1)
     if bias is not None:
-        dequantize_symmetric_with_bias(result, bias, scale, return_dtype, mm_output_shape)
+        dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape)
     else:
         dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
 
@@ -278,7 +278,7 @@ def conv_int8_matmul(
             result.append(torch._int_mm(input[i], weight[i]))
         result = torch.cat(result, dim=-1)
     if bias is not None:
-        result = dequantize_symmetric_with_bias(result, bias, scale, return_dtype, mm_output_shape)
+        result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape)
     else:
         result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
 
diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py
index e4717086e..b20c61818 100644
--- a/modules/sdnq/packed_int.py
+++ b/modules/sdnq/packed_int.py
@@ -6,11 +6,11 @@ import torch
 from .common import dtype_dict
 
 
-def pack_int_symetric(tensor: torch.ByteTensor, weights_dtype: str) -> torch.ByteTensor:
-    return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"]).sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
+def pack_int_symetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.ByteTensor:
+    return packed_int_function_dict[weights_dtype]["pack"](tensor.sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
 
 
-def unpack_int_symetric(packed_tensor: torch.CharTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.ByteTensor:
+def unpack_int_symetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.CharTensor:
     if dtype is None:
         dtype = dtype_dict[weights_dtype]["torch_dtype"]
     result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"])
@@ -19,9 +19,7 @@ def unpack_int_symetric(packed_tensor: torch.CharTensor, shape: torch.Size, weig
     return result
 
 
-def pack_uint7(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor:
     packed_tensor = tensor.contiguous().reshape(-1, 8)
     packed_tensor = torch.stack(
         (
@@ -38,9 +36,7 @@ def pack_uint7(tensor: torch.Tensor) -> torch.Tensor:
     return packed_tensor
 
 
-def pack_uint6(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor:
     packed_tensor = tensor.contiguous().reshape(-1, 4)
     packed_tensor = torch.stack(
         (
@@ -53,9 +49,7 @@ def pack_uint6(tensor: torch.Tensor) -> torch.Tensor:
     return packed_tensor
 
 
-def pack_uint5(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor:
     packed_tensor = tensor.contiguous().reshape(-1, 8)
     packed_tensor = torch.stack(
         (
@@ -82,17 +76,13 @@ def pack_uint5(tensor: torch.Tensor) -> torch.Tensor:
     return packed_tensor
 
 
-def pack_uint4(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+def pack_uint4(tensor: torch.ByteTensor) -> torch.ByteTensor:
     packed_tensor = tensor.contiguous().reshape(-1, 2)
     packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4))
     return packed_tensor
 
 
-def pack_uint3(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor:
     packed_tensor = tensor.contiguous().reshape(-1, 8)
     packed_tensor = torch.stack(
         (
@@ -117,9 +107,7 @@ def pack_uint3(tensor: torch.Tensor) -> torch.Tensor:
     return packed_tensor
 
 
-def pack_uint2(tensor: torch.Tensor) -> torch.Tensor:
-    if tensor.dtype != torch.uint8:
-        raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.")
+def pack_uint2(tensor: torch.ByteTensor) -> torch.ByteTensor:
     packed_tensor = tensor.contiguous().reshape(-1, 4)
     packed_tensor = torch.bitwise_or(
         torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)),
@@ -128,7 +116,7 @@ def pack_uint2(tensor: torch.Tensor) -> torch.Tensor:
     return packed_tensor
 
 
-def unpack_uint7(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
     result = torch.stack(
         (
             torch.bitwise_and(packed_tensor[:, 0], 127),
@@ -163,7 +151,7 @@ def unpack_uint7(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor
     return result
 
 
-def unpack_uint6(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
     result = torch.stack(
         (
             torch.bitwise_and(packed_tensor[:, 0], 63),
@@ -182,7 +170,7 @@ def unpack_uint6(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor
     return result
 
 
-def unpack_uint5(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+def unpack_uint5(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
     result = torch.stack(
         (
             torch.bitwise_and(packed_tensor[:, 0], 31),
@@ -211,12 +199,12 @@ def unpack_uint5(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor
     return result
 
 
-def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
     result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
     return result
 
 
-def unpack_uint3(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
     result = torch.stack(
         (
             torch.bitwise_and(packed_tensor[:, 0], 7),
@@ -239,7 +227,7 @@ def unpack_uint3(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor
     return result
 
 
-def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
     result = torch.stack(
         (
             torch.bitwise_and(packed_tensor, 3),
diff --git a/modules/shared.py b/modules/shared.py
index 2231e0484..d87d5b252 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -66,6 +66,7 @@ dir_timestamps = {}
 dir_cache = {}
 max_workers = 8
 default_hfcache_dir = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub')
+sdnq_quant_modes = ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"]
 
 
 class Backend(Enum):
@@ -501,7 +502,7 @@ options_templates.update(options_section(('backends', "Backend Settings"), {
     "olive_cache_optimized": OptionInfo(True, 'Olive cache optimized models'),
 
     "ipex_sep": OptionInfo("<h2>IPEX</h2>", "", gr.HTML, {"visible": devices.backend == "ipex"}),
-    "ipex_optimize": OptionInfo([], "IPEX Optimize", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE", "Upscaler"], "visible": devices.backend == "ipex"}),
+    "ipex_optimize": OptionInfo([], "IPEX Optimize", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE", "Upscaler"], "visible": devices.backend == "ipex"}),
 
     "openvino_sep": OptionInfo("<h2>OpenVINO</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
     "openvino_devices": OptionInfo([], "OpenVINO devices to use", gr.CheckboxGroup, {"choices": get_openvino_device_list() if cmd_opts.use_openvino else [], "visible": cmd_opts.use_openvino}), # pylint: disable=E0606
@@ -516,55 +517,55 @@ options_templates.update(options_section(('backends', "Backend Settings"), {
 
 options_templates.update(options_section(("quantization", "Quantization Settings"), {
     "sdnq_quantize_sep": OptionInfo("<h2>SDNQ: SD.Next Quantization</h2>", "", gr.HTML),
-    "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
-    "sdnq_quantize_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ["pre", "post"], "visible": native}),
-    "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}),
-    "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ["default", "int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}),
+    "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"], "visible": native}),
+    "sdnq_quantize_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"], "visible": native}),
+    "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes, "visible": native}),
+    "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['default'] + sdnq_quant_modes, "visible": native}),
     "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}),
-    "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}),
+    "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox, {"visible": native}),
     "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox, {"visible": native}),
-    "sdnq_use_quantized_matmul": OptionInfo(False, "Use Quantized MatMul", gr.Checkbox, {"visible": native}),
-    "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use Quantized MatMul with convolutional layers", gr.Checkbox, {"visible": native}),
-    "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize with the GPU", gr.Checkbox, {"visible": native}),
+    "sdnq_use_quantized_matmul": OptionInfo(False, "Use quantized MatMul", gr.Checkbox, {"visible": native}),
+    "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use quantized MatMul with convolutional layers", gr.Checkbox, {"visible": native}),
+    "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize using GPU", gr.Checkbox, {"visible": native}),
     "sdnq_dequantize_fp32": OptionInfo(False, "Dequantize using full precision", gr.Checkbox, {"visible": native}),
     "sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
 
     "bnb_quantization_sep": OptionInfo("<h2>BitsAndBytes</h2>", "", gr.HTML),
-    "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
+    "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "VAE"], "visible": native}),
     "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ["nf4", "fp8", "fp4"], "visible": native}),
     "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}),
 
     "quanto_quantization_sep": OptionInfo("<h2>Optimum Quanto</h2>", "", gr.HTML),
-    "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
+    "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM"], "visible": native}),
     "quanto_quantization_type": OptionInfo("int8", "Quantization weights type", gr.Dropdown, {"choices": ["float8", "int8", "int4", "int2"], "visible": native}),
 
     "optimum_quanto_sep": OptionInfo("<h2>Optimum Quanto: post-load</h2>", "", gr.HTML),
-    "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
+    "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "Control", "VAE"], "visible": native}),
     "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ["qint8", "qfloat8_e4m3fn", "qfloat8_e5m2", "qint4", "qint2"], "visible": native}),
     "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ["none", "qint8", "qfloat8_e4m3fn", "qfloat8_e5m2"], "visible": native}),
     "optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
 
     "torchao_sep": OptionInfo("<h2>TorchAO</h2>", "", gr.HTML),
-    "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "visible": native}),
-    "torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ["pre", "post"], "visible": native}),
+    "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"], "visible": native}),
+    "torchao_quantization_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"], "visible": native}),
     "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ["int4_weight_only", "int8_dynamic_activation_int4_weight", "int8_weight_only", "int8_dynamic_activation_int8_weight", "float8_weight_only", "float8_dynamic_activation_float8_weight", "float8_static_activation_float8_weight"], "visible": native}),
 
     "layerwise_quantization_sep": OptionInfo("<h2>Layerwise Casting</h2>", "", gr.HTML),
-    "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "TE"], "visible": native}),
+    "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"], "visible": native}),
     "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"], "visible": native}),
     "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox, {"visible": native}),
 
     "nunchaku_sep": OptionInfo("<h2>Nunchaku Engine</h2>", "", gr.HTML),
-    "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
+    "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"], "visible": native}),
     "nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox, {"visible": native}),
     "nunchaku_offload": OptionInfo(False, "Nunchaku offloading", gr.Checkbox, {"visible": native}),
 
     "nncf_compress_sep": OptionInfo("<h2>NNCF: Neural Network Compression Framework</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
-    "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": cmd_opts.use_openvino}),
+    "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}),
     "nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ["INT8", "INT4_ASYM", "INT8_SYM", "INT4_SYM", "NF4"], "visible": cmd_opts.use_openvino}),
     "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}),
     "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": cmd_opts.use_openvino}),
-    "nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE"], "visible": cmd_opts.use_openvino}),
+    "nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}),
     "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ["INT8", "FP8_E4M3", "FP8_E5M2"], "visible": cmd_opts.use_openvino}),
 }))
 
@@ -649,7 +650,7 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
 
 options_templates.update(options_section(('compile', "Model Compile"), {
     "cuda_compile_sep": OptionInfo("<h2>Model Compile</h2>", "", gr.HTML),
-    "cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE", "Upscaler"]}),
+    "cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE", "Upscaler"]}),
     "cuda_compile_backend": OptionInfo("none" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'migraphx', 'ipex', 'onediff', 'stable-fast', 'deep-cache', 'olive-ai', 'openvino_fx']}),
     "cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs']}),
     "cuda_compile_fullgraph": OptionInfo(True if not cmd_opts.use_openvino else False, "Model compile fullgraph"),
@@ -894,10 +895,11 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), {
 }))
 
 options_templates.update(options_section(('control', "Control Options"), {
-    "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
-    "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options"),
-    "control_move_processor": OptionInfo(False, "Processor move to CPU after use"),
-    "control_unload_processor": OptionInfo(False, "Processor unload after use"),
+    "control_hires": OptionInfo(False, "Use control during hires", gr.Checkbox, {"visible": False}),
+    "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1, "visible": False}),
+    "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options", gr.Textbox, {"visible": False}),
+    "control_move_processor": OptionInfo(False, "Processor move to CPU after use", gr.Checkbox, {"visible": False}),
+    "control_unload_processor": OptionInfo(False, "Processor unload after use", gr.Checkbox, {"visible": False}),
 }))
 
 options_templates.update(options_section(('interrogate', "Interrogate"), {
diff --git a/modules/shared_items.py b/modules/shared_items.py
index 9a09b8c7f..214d05801 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -27,6 +27,7 @@ pipelines = {
     'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None),
     'FLUX': getattr(diffusers, 'FluxPipeline', None),
     'FLEX': getattr(diffusers, 'AutoPipelineForText2Image', None),
+    'Chroma': getattr(diffusers, 'ChromaPipeline', None),
     'Sana': getattr(diffusers, 'SanaPipeline', None),
     'Lumina-Next': getattr(diffusers, 'LuminaText2ImgPipeline', None),
     'Lumina 2': getattr(diffusers, 'Lumina2Pipeline', None),
@@ -41,10 +42,11 @@ pipelines = {
     'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None),
     'Amused': getattr(diffusers, 'AmusedPipeline', None),
     'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None),
+    'OmniGenPipeline': getattr(diffusers, 'OmniGenPipeline', None),
+    'Cosmos': getattr(diffusers, 'Cosmos2TextToImagePipeline', None),
 
     # dynamically imported and redefined later
     'Meissonic': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
-    'OmniGenPipeline': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
     'InstaFlow': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
     'SegMoE': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
 }
diff --git a/modules/teacache/__init__.py b/modules/teacache/__init__.py
index 2fb81b7c9..56a261241 100644
--- a/modules/teacache/__init__.py
+++ b/modules/teacache/__init__.py
@@ -4,9 +4,10 @@ from .teacache_lumina2 import teacache_lumina2_forward
 from .teacache_ltx import teacache_ltx_forward
 from .teacache_mochi import teacache_mochi_forward
 from .teacache_cogvideox import teacache_cog_forward
+from .teacache_chroma import teacache_chroma_forward
 
 
-supported_models = ['Flux', 'CogVideoX', 'Mochi', 'LTX', 'HiDream', 'Lumina2']
+supported_models = ['Flux', 'Chroma', 'CogVideoX', 'Mochi', 'LTX', 'HiDream', 'Lumina2']
 
 
 def apply_teacache(p):
diff --git a/modules/teacache/teacache_chroma.py b/modules/teacache/teacache_chroma.py
new file mode 100644
index 000000000..102ae7f9b
--- /dev/null
+++ b/modules/teacache/teacache_chroma.py
@@ -0,0 +1,332 @@
+from typing import Any, Dict, Optional, Union
+import torch
+import numpy as np
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def teacache_chroma_forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+    """
+    The [`ChromaTransformer2DModel`] forward method.
+    Args:
+        hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+            Input `hidden_states`.
+        encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+            Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+        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.
+    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."
+            )
+
+    hidden_states = self.x_embedder(hidden_states)
+
+    timestep = timestep.to(hidden_states.dtype) * 1000
+
+    input_vec = self.time_text_embed(timestep)
+    pooled_temb = self.distilled_guidance_layer(input_vec)
+
+    encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+    if txt_ids.ndim == 3:
+        logger.warning(
+            "Passing `txt_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        txt_ids = txt_ids[0]
+    if img_ids.ndim == 3:
+        logger.warning(
+            "Passing `img_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        img_ids = img_ids[0]
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    image_rotary_emb = self.pos_embed(ids)
+
+    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 = self.encoder_hid_proj(ip_adapter_image_embeds)
+        joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+    if self.enable_teacache:
+        inp = hidden_states.clone()
+        input_vec_ = input_vec.clone()
+        modulated_inp, _gate_msa, _shift_mlp, _scale_mlp, _gate_mlp = self.transformer_blocks[0].norm1(inp, emb=input_vec_)
+        if self.cnt == 0 or self.cnt == self.num_steps-1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else:
+            coefficients = [4.98651651e+02, -2.83781631e+02,  5.58554382e+01, -3.82021401e+00, 2.64230861e-01]
+            rescale_func = np.poly1d(coefficients)
+            self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = modulated_inp
+        self.cnt += 1
+        if self.cnt == self.num_steps:
+            self.cnt = 0
+
+    if self.enable_teacache:
+        if not should_calc:
+            hidden_states += self.previous_residual
+        else:
+            ori_hidden_states = hidden_states.clone()
+            for index_block, block in enumerate(self.transformer_blocks):
+                img_offset = 3 * len(self.single_transformer_blocks)
+                txt_offset = img_offset + 6 * len(self.transformer_blocks)
+                img_modulation = img_offset + 6 * index_block
+                text_modulation = txt_offset + 6 * index_block
+                temb = torch.cat(
+                    (
+                        pooled_temb[:, img_modulation : img_modulation + 6],
+                        pooled_temb[:, text_modulation : text_modulation + 6],
+                    ),
+                    dim=1,
+                )
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                    def create_custom_forward4(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward4(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        image_rotary_emb,
+                        attention_mask,
+                        **ckpt_kwargs,
+                    )
+
+                else:
+                    encoder_hidden_states, hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        attention_mask=attention_mask,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+
+                # controlnet residual
+                if controlnet_block_samples is not None:
+                    interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                    interval_control = int(np.ceil(interval_control))
+                    # For Xlabs ControlNet.
+                    if controlnet_blocks_repeat:
+                        hidden_states = (
+                            hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                        )
+                    else:
+                        hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+            for index_block, block in enumerate(self.single_transformer_blocks):
+                start_idx = 3 * index_block
+                temb = pooled_temb[:, start_idx : start_idx + 3]
+
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                    def create_custom_forward2(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward2(block),
+                        hidden_states,
+                        temb,
+                        image_rotary_emb,
+                        **ckpt_kwargs,
+                    )
+
+                else:
+                    hidden_states = block(
+                        hidden_states=hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+
+                # controlnet residual
+                if controlnet_single_block_samples is not None:
+                    interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                    interval_control = int(np.ceil(interval_control))
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                        hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                        + controlnet_single_block_samples[index_block // interval_control]
+                    )
+
+            hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+            self.previous_residual = hidden_states - ori_hidden_states
+    else:
+        for index_block, block in enumerate(self.transformer_blocks):
+            img_offset = 3 * len(self.single_transformer_blocks)
+            txt_offset = img_offset + 6 * len(self.transformer_blocks)
+            img_modulation = img_offset + 6 * index_block
+            text_modulation = txt_offset + 6 * index_block
+            temb = torch.cat(
+                (
+                    pooled_temb[:, img_modulation : img_modulation + 6],
+                    pooled_temb[:, text_modulation : text_modulation + 6],
+                ),
+                dim=1,
+            )
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward1(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward1(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    attention_mask=attention_mask,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    attention_mask=attention_mask,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            start_idx = 3 * index_block
+            temb = pooled_temb[:, start_idx : start_idx + 3]
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward3(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward3(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    attention_mask=attention_mask,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    attention_mask=attention_mask,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+    temb = pooled_temb[:, -2:]
+    hidden_states = self.norm_out(hidden_states, temb)
+    output = self.proj_out(hidden_states)
+
+    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)
\ No newline at end of file
diff --git a/modules/timer.py b/modules/timer.py
index 69107e605..59c6a1de3 100644
--- a/modules/timer.py
+++ b/modules/timer.py
@@ -44,8 +44,7 @@ class Timer:
     def summary(self, min_time=default_min_time, total=True):
         if self.profile:
             min_time = -1
-        if self.total <= 0:
-            self.total = sum(self.records.values())
+        self.total = sum(self.records.values())
         res = f"total={self.total:.2f} " if total else ''
         additions = [x for x in self.records.items() if x[1] >= min_time]
         additions = sorted(additions, key=lambda x: x[1], reverse=True)
@@ -60,6 +59,8 @@ class Timer:
     def dct(self, min_time=default_min_time):
         if self.profile:
             res = {k: round(v, 4) for k, v in self.records.items()}
+        self.total = sum(self.records.values())
+        self.records['total'] = self.total
         res = {k: round(v, 2) for k, v in self.records.items() if v >= min_time}
         res = {k: v for k, v in sorted(res.items(), key=lambda x: x[1], reverse=True)} # noqa: C416 # pylint: disable=unnecessary-comprehension
         return res
diff --git a/modules/ui_common.py b/modules/ui_common.py
index d0696461a..84cae441d 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -54,7 +54,7 @@ def infotext_to_html(text):
     return code
 
 
-def delete_files(js_data, files, _html_info, index):
+def delete_files(js_data, files, all_files, index):
     try:
         data = json.loads(js_data)
     except Exception:
@@ -63,25 +63,26 @@ def delete_files(js_data, files, _html_info, index):
     if index > -1 and shared.opts.save_selected_only and (index >= data['index_of_first_image']):
         files = [files[index]]
         start_index = index
-        filenames = []
-    filenames = []
-    fullfns = []
+    deleted = []
+    all_files = [f.split('/file=')[1] if 'file=' in f else f for f in all_files] if isinstance(all_files, list) else []
     for _image_index, filedata in enumerate(files, start_index):
-        if 'name' in filedata and os.path.isfile(filedata['name']):
-            fullfn = filedata['name']
-            filenames.append(os.path.basename(fullfn))
-            try:
-                os.remove(fullfn)
-                base, _ext = os.path.splitext(fullfn)
-                desc = f'{base}.txt'
-                if os.path.exists(desc):
-                    os.remove(desc)
-                fullfns.append(fullfn)
-                shared.log.info(f"Deleting image: {fullfn}")
-            except Exception as e:
-                shared.log.error(f'Error deleting file: {fullfn} {e}')
-    files = [image for image in files if image['name'] not in fullfns]
-    return files, plaintext_to_html(f"Deleted: {filenames[0] if len(filenames) > 0 else 'none'}")
+        try:
+            fn = filedata['name']
+            if os.path.isfile(fn):
+                deleted.append(fn)
+                os.remove(fn)
+                if fn in all_files:
+                    all_files.remove(fn)
+                shared.log.info(f'Delete: image="{fn}"')
+            base, _ext = os.path.splitext(fn)
+            desc = f'{base}.txt'
+            if os.path.exists(desc):
+                os.remove(desc)
+                shared.log.info(f'Delete: text="{fn}"')
+        except Exception as e:
+            shared.log.error(f'Delete: image="{fn}" {e}')
+    deleted = ', '.join(deleted) if len(deleted) > 0 else 'none'
+    return all_files, plaintext_to_html(f"Deleted: {deleted}")
 
 
 def save_files(js_data, files, html_info, index):
@@ -296,8 +297,8 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None, transfe
                     inputs=[generation_info, result_gallery, html_info, html_info],
                     outputs=[download_files, html_log],
                 )
-                delete.click(fn=call_queue.wrap_gradio_call(delete_files),show_progress=False,
-                    _js="(x, y, z, i) => [x, y, z, selected_gallery_index()]",
+                delete.click(fn=call_queue.wrap_gradio_call(delete_files), show_progress=False,
+                    _js="(x, y, i, j) => [x, y, ...selected_gallery_files()]",
                     inputs=[generation_info, result_gallery, html_info, html_info],
                     outputs=[result_gallery, html_log],
                 )
diff --git a/modules/ui_control.py b/modules/ui_control.py
index 94ad27eee..4b53a76bf 100644
--- a/modules/ui_control.py
+++ b/modules/ui_control.py
@@ -467,9 +467,31 @@ def create_ui(_blocks: gr.Blocks=None):
                             if i == 0:
                                 units[-1].enabled = True # enable first unit in group
 
-                with gr.Accordion('Processor settings', open=False, elem_classes=['control-settings']) as _tab_settings:
+                with gr.Accordion('Control settings', open=False, elem_classes=['control-settings']) as _tab_settings:
                     with gr.Group(elem_classes=['processor-group']):
                         settings = []
+                        with gr.Accordion('Global', open=True, elem_classes=['processor-settings']):
+                            control_hires = gr.Checkbox(label="Use control during hires", value=shared.opts.control_hires, elem_id='control_hires')
+                            def set_control_hires(value):
+                                shared.opts.control_active = value
+                            control_hires.change(fn=set_control_hires, inputs=[control_hires], outputs=[])
+                            control_max_units = gr.Slider(label="Maximum units", minimum=1, maximum=10, step=1, value=shared.opts.control_max_units, elem_id='control_max_units')
+                            def set_control_max_units(value):
+                                shared.opts.control_max_units = value
+                            control_max_units.change(fn=set_control_max_units, inputs=[control_max_units], outputs=[])
+                            control_tiles = gr.Textbox(label="Tiling options", value=shared.opts.control_tiles, elem_id='control_tiles')
+                            def set_control_tiles(value):
+                                shared.opts.control_tiles = value
+                            control_tiles.change(fn=set_control_tiles, inputs=[control_tiles], outputs=[])
+                            control_move_processor = gr.Checkbox(label="Move processor to CPU after use", value=shared.opts.control_move_processor, elem_id='control_move_processor')
+                            def set_control_move_processor(value):
+                                shared.opts.control_move_processor = value
+                            control_move_processor.change(fn=set_control_move_processor, inputs=[control_move_processor], outputs=[])
+                            control_unload_processor = gr.Checkbox(label="Unload processor after use", value=shared.opts.control_unload_processor, elem_id='control_unload_processor')
+                            def set_control_unload_processor(value):
+                                shared.opts.control_unload_processor = value
+                            control_unload_processor.change(fn=set_control_unload_processor, inputs=[control_unload_processor], outputs=[])
+
                         with gr.Accordion('HED', open=True, elem_classes=['processor-settings']):
                             settings.append(gr.Checkbox(label="Scribble", value=False))
                         with gr.Accordion('Midas depth', open=True, elem_classes=['processor-settings']):
diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index d2f8e4576..ca4c5c855 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -941,9 +941,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         from modules.processing_info import get_last_args
         params, text = get_last_args()
         if (not params) or (not text) or (len(text) == 0):
-            filename = os.path.join(paths.data_path, "params.txt")
-            if os.path.exists(filename):
-                with open(filename, "r", encoding="utf8") as file:
+            if os.path.exists(paths.params_path):
+                with open(paths.params_path, "r", encoding="utf8") as file:
                     text = file.read()
             else:
                 text = ''
@@ -960,9 +959,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         from modules.processing_info import get_last_args
         params, text = get_last_args()
         if (not params) or (not text) or (len(text) == 0):
-            fn = os.path.join(paths.data_path, "params.txt")
-            if os.path.exists(fn):
-                with open(fn, "r", encoding="utf8") as file:
+            if os.path.exists(paths.params_path):
+                with open(paths.params_path, "r", encoding="utf8") as file:
                     text = file.read()
             else:
                 text = ''
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 704399158..82ed79298 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -37,7 +37,14 @@ def create_ui():
                     model = modelstats.analyze()
                     desc = f"Model: {model.name}<br>Type: {model.type}<br>Class: {model.cls}<br>Size: {model.size} bytes<br>Modified: {model.mtime}<br>"
                     meta = model.meta
-                    components = [(m.name, m.cls, m.device, m.dtype, m.params, m.modules, str(m.config)) for m in model.modules]
+                    components = []
+                    for m in model.modules:
+                        try:
+                            component = (m.name, m.cls, str(m.device), str(m.dtype), m.params, m.modules, str(m.config))
+                            components.append(component)
+                        except Exception:
+                            component = (m.name, m.cls, str(m.device), str(m.dtype), m.params, m.modules, '')
+                            components.append(component)
                     return [desc, components, meta]
 
                 with gr.Row():
diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py
index 33813df82..2a959f0d3 100644
--- a/modules/video_models/video_load.py
+++ b/modules/video_models/video_load.py
@@ -1,4 +1,5 @@
 import os
+import copy
 import time
 from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devices, sd_hijack_te
 from modules.video_models import models_def, video_utils, video_vae, video_overrides, video_cache
@@ -37,7 +38,7 @@ def load_model(selected: models_def.Model):
 
     # transformer
     try:
-        quant_args = model_quant.create_config(module='Video')
+        quant_args = model_quant.create_config(module='Model')
         debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" folder="{selected.dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
         transformer = selected.dit_cls.from_pretrained(
             pretrained_model_name_or_path=selected.dit or selected.repo,
@@ -70,6 +71,9 @@ def load_model(selected: models_def.Model):
         errors.display(e, 'video')
 
     t1 = time.time()
+    if shared.sd_model.__class__.__name__.startswith("LTX"):
+        shared.sd_model.scheduler.config.use_dynamic_shifting = False
+    shared.sd_model.default_scheduler = copy.deepcopy(shared.sd_model.scheduler) if hasattr(shared.sd_model, "scheduler") else None
     shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(selected.repo)
     shared.sd_model.sd_model_hash = None
     sd_models.set_diffuser_options(shared.sd_model)
diff --git a/requirements.txt b/requirements.txt
index f5eb92014..0eea8fbac 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 # required for python 3.12
 setuptools==69.5.1
+wheel
 
 # standard
 patch-ng
@@ -41,11 +42,11 @@ torchsde==0.2.6
 antlr4-python3-runtime==4.9.3
 requests==2.32.3
 tqdm==4.67.1
-accelerate==1.7.0
+accelerate==1.8.1
 opencv-contrib-python-headless==4.9.0.80
 einops==0.4.1
 gradio==3.43.2
-huggingface_hub==0.33.0
+huggingface_hub==0.33.1
 numexpr==2.10.2
 numpy==1.26.4
 pandas==2.3.0
@@ -53,10 +54,10 @@ numba==0.61.2
 protobuf==4.25.3
 pytorch_lightning==1.9.4
 tokenizers==0.21.1
-transformers==4.52.4
+transformers==4.53.0
 urllib3==1.26.19
 Pillow==10.4.0
-timm==0.9.16
+timm==1.0.16
 pydantic==1.10.21
 pyparsing==3.1.4
 typing-extensions==4.12.2
diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py
index e1b9c5502..26981f812 100644
--- a/scripts/prompt_enhance.py
+++ b/scripts/prompt_enhance.py
@@ -367,6 +367,7 @@ class Script(scripts.Script):
                 clean_up_tokenization_spaces=True,
             )
         except Exception as e:
+            outputs = None
             shared.log.error(f'Prompt enhance generate: {e}')
             errors.display(e, 'Prompt enhance')
             self.busy = False
diff --git a/scripts/prompt_matrix.py b/scripts/prompt_matrix.py
index 5babf20e7..37cc2795c 100644
--- a/scripts/prompt_matrix.py
+++ b/scripts/prompt_matrix.py
@@ -21,7 +21,7 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell):
         for ix, x in enumerate(xs):
             state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"
 
-            processed = cell(x, y)
+            processed, t = cell(x, y)
             if first_processed is None:
                 first_processed = processed
 
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 240376b3c..c79ce6909 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -1,6 +1,7 @@
 # xyz grid that shows as selectable script
 import os
 import csv
+import time
 import random
 from collections import namedtuple
 from copy import copy
@@ -303,7 +304,7 @@ class Script(scripts.Script):
 
         def cell(x, y, z, ix, iy, iz):
             if shared.state.interrupted:
-                return processing.Processed(p, [], p.seed, "")
+                return processing.Processed(p, [], p.seed, ""), 0
             p.xyz = True
             pc = copy(p)
             pc.override_settings_restore_afterwards = False
@@ -311,6 +312,8 @@ class Script(scripts.Script):
             x_opt.apply(pc, x, xs)
             y_opt.apply(pc, y, ys)
             z_opt.apply(pc, z, zs)
+
+            t0 = time.time()
             try:
                 processed = processing.process_images(pc)
             except Exception as e:
@@ -341,7 +344,8 @@ class Script(scripts.Script):
                         pc.extra_generation_params["Fixed Z Values"] = ", ".join([str(z) for z in zs])
                 grid_text = f'{len(zs)}x{len(xs)}x{len(ys)}' if len(zs) > 0 else f'{len(xs)}x{len(ys)}'
                 grid_infotext[0] = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds, grid=grid_text)
-            return processed
+            t1 = time.time()
+            return processed, t1-t0
 
         with SharedSettingsStackHelper():
             processed = draw_xyz_grid(
diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py
index 0d3d6dd97..289059afb 100644
--- a/scripts/xyz_grid_classes.py
+++ b/scripts/xyz_grid_classes.py
@@ -1,4 +1,4 @@
-from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
+from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, apply_sdnq_quant,  apply_sdnq_quant_te, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
 from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet
 
 
@@ -58,6 +58,7 @@ class SharedSettingsStackHelper(object):
     extra_networks_default_multiplier = None
     disable_apply_metadata = None
     disable_apply_params = None
+    sdnq_quant_mode = None
 
     def __enter__(self):
         # Save overridden settings so they can be restored later
@@ -89,6 +90,8 @@ class SharedSettingsStackHelper(object):
         self.teacache_thresh = shared.opts.teacache_thresh
         self.disable_apply_metadata = shared.opts.disable_apply_metadata
         self.disable_apply_params = shared.opts.disable_apply_params
+        self.sdnq_quant_mode = shared.opts.sdnq_quantize_weights_mode
+
         shared.opts.data["disable_apply_metadata"] = []
         shared.opts.data["disable_apply_params"] = ''
 
@@ -135,6 +138,9 @@ class SharedSettingsStackHelper(object):
         if self.sd_unet != shared.opts.sd_unet:
             shared.opts.data["sd_unet"] = self.sd_unet
             sd_unet.load_unet(shared.sd_model)
+        if self.sdnq_quant_mode != shared.opts.sdnq_quantize_weights_mode:
+            shared.opts.data["sdnq_quantize_weights_mode"] = self.sdnq_quant_mode
+            sd_models.reload_model_weights(op='model')
 
 
 axis_options = [
@@ -193,6 +199,8 @@ axis_options = [
     AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]),
     AxisOption("[Postprocess] Detailer", str, apply_detailer, fmt=format_value_add_label),
     AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")),
+    AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
+    AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
     AxisOption("[HDR] Mode", int, apply_field("hdr_mode")),
     AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")),
     AxisOption("[HDR] Color", float, apply_field("hdr_color")),
diff --git a/scripts/xyz_grid_draw.py b/scripts/xyz_grid_draw.py
index d654a2c7f..646a3bd52 100644
--- a/scripts/xyz_grid_draw.py
+++ b/scripts/xyz_grid_draw.py
@@ -22,9 +22,9 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
         def index(ix, iy, iz):
             return ix + iy * len(xs) + iz * len(xs) * len(ys)
 
-        p0 = time.time()
-        processed: processing.Processed = cell(x, y, z, ix, iy, iz)
-        p1 = time.time()
+        res = cell(x, y, z, ix, iy, iz)
+        processed: processing.Processed = res[0] if isinstance(res, tuple) else res
+        elapsed = res[1] if isinstance(res, tuple) else 0
         if processed_result is None:
             processed_result = copy(processed)
             if processed_result is None:
@@ -48,13 +48,13 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
                 if len(z_labels[iz]) > 0:
                     overlay_text += f'{z_labels[iz]}\n'
             if include_time:
-                overlay_text += f'Time: {p1 - p0:.2f}'
+                overlay_text += f'Time: {elapsed:.2f}'
             if len(overlay_text) > 0:
                 processed_result.images[idx] = images.draw_overlay(processed_result.images[idx], overlay_text)
             processed_result.all_prompts[idx] = processed.prompt
             processed_result.all_seeds[idx] = processed.seed
             processed_result.infotexts[idx] = processed.infotexts[0]
-            processed_result.time[idx] = round(p1 - p0, 2)
+            processed_result.time[idx] = round(elapsed, 2)
         else:
             cell_mode = "P"
             cell_size = (processed_result.width, processed_result.height)
diff --git a/scripts/xyz_grid_on.py b/scripts/xyz_grid_on.py
index f40a1731b..5e070ce8b 100644
--- a/scripts/xyz_grid_on.py
+++ b/scripts/xyz_grid_on.py
@@ -1,6 +1,7 @@
 # xyz grid that shows up as alwayson script
 import os
 import csv
+import time
 import random
 from collections import namedtuple
 from copy import copy
@@ -322,7 +323,7 @@ class Script(scripts.Script):
 
         def cell(x, y, z, ix, iy, iz):
             if shared.state.interrupted:
-                return processing.Processed(p, [], p.seed, "")
+                return processing.Processed(p, [], p.seed, ""), 0
             p.xyz = True
             pc = copy(p)
             pc.override_settings_restore_afterwards = False
@@ -331,6 +332,7 @@ class Script(scripts.Script):
             y_opt.apply(pc, y, ys)
             z_opt.apply(pc, z, zs)
 
+            t0 = time.time()
             try:
                 processed = processing.process_images(pc)
             except Exception as e:
@@ -362,7 +364,8 @@ class Script(scripts.Script):
                         pc.extra_generation_params["Fixed Z Values"] = ", ".join([str(z) for z in zs])
                 info = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds, grid=f'{len(zs)}x{len(xs)}x{len(ys)}')
                 grid_infotext.insert(0, info)
-            return processed
+            t1 = time.time()
+            return processed, t1-t0
 
         with SharedSettingsStackHelper():
             processed: processing.Processed = draw_xyz_grid(
diff --git a/scripts/xyz_grid_shared.py b/scripts/xyz_grid_shared.py
index 2d594bac7..64f179075 100644
--- a/scripts/xyz_grid_shared.py
+++ b/scripts/xyz_grid_shared.py
@@ -147,6 +147,20 @@ def confirm_samplers(p, xs):
             shared.log.warning(f"XYZ grid: unknown sampler: {x}")
 
 
+def apply_sdnq_quant(p, x, xs):
+    shared.opts.sdnq_quantize_weights_mode = x
+    sd_models.unload_model_weights(op='model')
+    sd_models.reload_model_weights()
+    shared.log.debug(f'XYZ grid apply sdnq quant: mode="{x}"')
+
+
+def apply_sdnq_quant_te(p, x, xs):
+    shared.opts.sdnq_quantize_weights_mode_te = x
+    sd_models.unload_model_weights(op='model')
+    sd_models.reload_model_weights()
+    shared.log.debug(f'XYZ grid apply sdnq quant te: mode="{x}"')
+
+
 def apply_checkpoint(p, x, xs):
     if x == shared.opts.sd_model_checkpoint:
         return
diff --git a/wiki b/wiki
index 2ca67daca..45c389a37 160000
--- a/wiki
+++ b/wiki
@@ -1 +1 @@
-Subproject commit 2ca67dacaccb7ef4c0595b19407fc5a93167008b
+Subproject commit 45c389a3747cb1caceb4035535041c400c29cad2