From 24194201cf11d1b3270dca8fec174e2b90e0b44c Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Sat, 14 Jun 2025 19:55:43 +0300
Subject: [PATCH] Fix OmniGen

---
 CHANGELOG.md                   | 1 +
 modules/omnigen/transformer.py | 5 +++++
 modules/processing_vae.py      | 4 ++--
 modules/sd_models_utils.py     | 2 ++
 4 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index abe785428..d3ec2c882 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -33,6 +33,7 @@
 
 - **Fixes**  
   - Meissonic with multiple generators  
+  - OmniGen with new transformers  
   - Invalid attention processors  
   - PixArt Sigma Small and Large loading  
   - TAESD previews with PixArt and Lumina 2  
diff --git a/modules/omnigen/transformer.py b/modules/omnigen/transformer.py
index f3bcdb15a..d166309ca 100644
--- a/modules/omnigen/transformer.py
+++ b/modules/omnigen/transformer.py
@@ -99,6 +99,9 @@ class Phi3Transformer(Phi3Model):
 
         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
@@ -118,6 +121,7 @@ class Phi3Transformer(Phi3Model):
                     output_attentions,
                     use_cache,
                     cache_position,
+                    position_embeddings,
                 )
             else:
                 layer_outputs = decoder_layer(
@@ -128,6 +132,7 @@ class Phi3Transformer(Phi3Model):
                     output_attentions=output_attentions,
                     use_cache=use_cache,
                     cache_position=cache_position,
+                    position_embeddings=position_embeddings,
                 )
 
             hidden_states = layer_outputs[0]
diff --git a/modules/processing_vae.py b/modules/processing_vae.py
index 290c0d489..36d5d0dda 100644
--- a/modules/processing_vae.py
+++ b/modules/processing_vae.py
@@ -256,8 +256,8 @@ def vae_postprocess(tensor, model, output_type='np'):
                 if output_type == "pil":
                     images = model.numpy_to_pil(images)
             else:
-                import diffusers
-                model.image_processor = diffusers.image_processor.VaeImageProcessor()
+                from diffusers.image_processor import VaeImageProcessor
+                model.image_processor = VaeImageProcessor()
                 images = model.image_processor.postprocess(tensor, output_type=output_type)
         else:
             images = tensor if isinstance(tensor, list) or isinstance(tensor, np.ndarray) else [tensor]
diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py
index 0e1b72f4c..546297d25 100644
--- a/modules/sd_models_utils.py
+++ b/modules/sd_models_utils.py
@@ -155,6 +155,8 @@ def apply_function_to_model(sd_model, function, options, op=None):
         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, 'decoder_pipe') and hasattr(sd_model, 'decoder'):