Merge pull request #4009 from vladmandic/dev

merge dev
2026-09-09 06:18:43 +02:00 · 2025-06-30 13:13:36 -04:00
parent 72eb013294 d6db6ff349
commit 0d7c025af1
130 changed files with 6547 additions and 1904 deletions
@@ -27,6 +27,7 @@ ignore-paths=/usr/lib/.*$,
             modules/meissonic,
             modules/mod,
             modules/omnigen,
+             modules/omnigen2,
             modules/onnx_impl,
             modules/pag,
             modules/pixelsmith,
@@ -20,6 +20,7 @@ exclude = [
    "modules/meissonic",
    "modules/mod",
    "modules/omnigen",
+    "modules/omnigen2",
    "modules/hidream",
    "modules/pag",
    "modules/pixelsmith",
@@ -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  
@@ -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
@@ -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

@@ -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"
+  ]
+}
@@ -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
+}
@@ -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
+}
@@ -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
+}
@@ -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
+  }
+}
@@ -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>"
+}
@@ -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
+}
@@ -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
+}
@@ -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
@@ -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 = {}
@@ -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).",
@@ -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")
@@ -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);
@@ -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]];
@@ -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)
@@ -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},
@@ -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()
@@ -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)
@@ -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()
@@ -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": []},
@@ -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()
@@ -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)
@@ -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:
@@ -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


@@ -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)
@@ -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}')
@@ -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
@@ -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,
@@ -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
@@ -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
@@ -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

@@ -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

@@ -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
@@ -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:
@@ -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
@@ -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)
@@ -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):
@@ -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
@@ -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
@@ -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:
@@ -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
@@ -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


@@ -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
@@ -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,
@@ -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
@@ -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:
@@ -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)

@@ -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}')

@@ -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}"')
@@ -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
@@ -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
@@ -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',
@@ -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)
@@ -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'):
@@ -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')
@@ -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
@@ -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'
@@ -1,4 +0,0 @@
-from .model import OmniGen
-from .processor import OmniGenProcessor
-from .scheduler import OmniGenScheduler
-from .pipeline import OmniGenPipeline
@@ -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
@@ -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
@@ -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
@@ -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
-
@@ -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,
-        )
-
@@ -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
@@ -0,0 +1,3 @@
+from transformers import Qwen2_5_VLForConditionalGeneration
+from .pipeline_omnigen2 import OmniGen2Pipeline
+from .models.transformers import OmniGen2Transformer2DModel
@@ -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
@@ -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
@@ -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)
@@ -0,0 +1,3 @@
+from .transformer_omnigen2 import OmniGen2Transformer2DModel
+
+__all__ = ["OmniGen2Transformer2DModel"]
@@ -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
@@ -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,
+        )
@@ -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)
@@ -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
@@ -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")
@@ -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=[])
@@ -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
@@ -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"):
@@ -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
@@ -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")
@@ -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()
@@ -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
@@ -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
@@ -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:
@@ -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()
@@ -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
@@ -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()

@@ -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:
--- a/Show More
+++ b/Show More