experimental xomni

Signed-off-by: Vladimir Mandic <mandic00@live.com>

.pylintrc

@@ -42,6 +42,7 @@ ignore-paths=/usr/lib/.*$,
              pipelines/meissonic,
              pipelines/omnigen2,
              pipelines/segmoe,
+             pipelines/xomni,
              scripts/consistory,
              scripts/ctrlx,
              scripts/daam,

.ruff.toml

@@ -22,6 +22,7 @@ exclude = [
     "pipelines/omnigen2",
     "pipelines/hdm",
     "pipelines/segmoe",
+    "pipelines/xomni",
 
     "scripts/lbm",
     "scripts/daam",

CHANGELOG.md

@@ -9,14 +9,16 @@
     updated version of Qwen Image Edit with improved image consistency  
   - [Qwen Image Pruning](https://huggingface.co/OPPOer/Qwen-Image-Pruning) and [Qwen Image Edit Pruning](https://huggingface.co/OPPOer/Qwen-Image-Edit-Pruning)  
     pruned versions of Qwen with 13B params instead of 20B, with some quality tradeoff  
-  - [HiDream E1.1](https://huggingface.co/HiDream-ai/HiDream-E1-1)  
-    updated version of E1 image editing model  
   - [Tencent FLUX.1 Dev SRPO](https://huggingface.co/tencent/SRPO)  
     SRPO is trained by Tencent with specific technique: directly aligning the full diffusion trajectory with fine-grained human preference  
   - [Nunchaku SDXL](https://huggingface.co/nunchaku-tech/nunchaku-sdxl) and [Nunchaku SDXL Turbo](https://huggingface.co/nunchaku-tech/nunchaku-sdxl-turbo)  
     impact of nunchaku engine on unet-based model such as sdxl is much less than on a dit-based models, but its still significantly faster than baseline  
     note that nunchaku optimized and prequantized unet is replacement for base unet, so its only applicable to base models, not any of finetunes  
     *how to use*: enable nunchaku in settings -> quantization and then load either sdxl-base or sdxl-base-turbo reference models  
+  - [X-Omni SFT](https://x-omni-team.github.io/)  
+    *experimental*: X-omni is a transformer-only discrete autoregressive image generative model trained with reinforcement learning  
+  - [HiDream E1.1](https://huggingface.co/HiDream-ai/HiDream-E1-1)  
+    *experimental*: updated version of E1 image editing model  
 - **Features**
   - [Qwen Image-Edit] multi-image editing
     requires qwen-image-edit-2509 or its variant as multi-image edits are not available in original qwen-image  
@@ -94,8 +96,9 @@
     also avoids creation of temporary files for each frame unless user wants to save them  
   - unified prompt enhance code across all video models  
   - add job state tracking for video generation  
-  - improve offloading for **ltx** and **wan**
-  - fix model selection in ltx tab
+  - fix quantization not being applied on load for some models  
+  - improve offloading for **ltx** and **wan**  
+  - fix model selection in **ltx** tab  
 - **Experimental**
   - `new` command line flag enables new `pydantic` and `albumentations` packages  
   - **modular pipelines**: enable in *settings -> model options*  

TODO.md

@@ -16,7 +16,6 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 ### Under Consideration
 
 - [Inf-DiT](https://github.com/zai-org/Inf-DiT)
-- [X-Omni](https://github.com/X-Omni-Team/X-Omni/blob/main/README.md)
 - [DiffSynth Studio](https://github.com/modelscope/DiffSynth-Studio)
 - [IPAdapter negative guidance](https://github.com/huggingface/diffusers/discussions/7167)  
 - [IPAdapter composition](https://huggingface.co/ostris/ip-composition-adapter)  

html/reference.json

@@ -501,6 +501,16 @@
     "date": "2025 June"
   },
 
+  "X-Omni SFT": {
+    "path": "X-Omni/X-Omni-SFT",
+    "desc": "X-Omni: Reinforcement learning makes discrete autoregressive image generative models great again",
+    "preview": "X-Omni--X-Omni-SFT.jpg",
+    "skip": true,
+    "size": 0,
+    "date": "2024 September",
+    "experimental": true
+  },
+
   "VectorSpaceLab 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.",

modules/modeldata.py

@@ -62,6 +62,8 @@ def get_model_type(pipe):
         model_type = 'qwen'
     elif 'NextStep' in name:
         model_type = 'nextstep'
+    elif 'X-Omni' in name:
+        model_type = 'x-omni'
     # video models
     elif "CogVideo" in name:
         model_type = 'cogvideo'

modules/sd_checkpoint.py

@@ -234,6 +234,8 @@ def get_closest_checkpoint_match(s: str) -> CheckpointInfo:
 
     # reference search
     ref = [(k, v) for k, v in shared.reference_models.items() if f"{v.get('path', '')}+{v.get('subfolder', '')}" == s]
+    if len(ref) == 0:
+        ref = [(k, v) for k, v in shared.reference_models.items() if v.get('path', '') == s]
     if ref and len(ref) > 0:
         _name, info = ref[0]
         checkpoint_info = CheckpointInfo(s)

modules/sd_detect.py

@@ -115,6 +115,8 @@ def guess_by_name(fn, current_guess):
         return 'Kandinsky 3.0'
     elif 'hunyuanimage' in fn.lower():
         return 'HunyuanImage'
+    elif 'x-omni' in fn.lower():
+        return 'X-Omni'
     elif 'sdxl-turbo' in fn.lower() or 'stable-diffusion-xl' in fn.lower():
         return 'Stable Diffusion XL'
     return current_guess

modules/sd_models.py

@@ -47,6 +47,7 @@ pipe_switch_task_exclude = [
     'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline',
     'StableDiffusionReferencePipeline',
     'StableDiffusionXLInstantIDPipeline',
+    'XOmniPipeline',
 ]
 i2i_pipes = [
     'LEditsPPPipelineStableDiffusion',
@@ -399,6 +400,10 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
             from pipelines.model_hyimage import load_hyimage
             sd_model = load_hyimage(checkpoint_info, diffusers_load_config) # pylint: disable=assignment-from-none
             allow_post_quant = False
+        elif model_type in ['X-Omni']:
+            from pipelines.model_xomni import load_xomni
+            sd_model = load_xomni(checkpoint_info, diffusers_load_config) # pylint: disable=assignment-from-none
+            allow_post_quant = False
     except Exception as e:
         shared.log.error(f'Load {op}: path="{checkpoint_info.path}" {e}')
         if debug_load:

modules/sd_offload.py

@@ -13,7 +13,7 @@ from modules.timer import process as process_timer
 
 debug = os.environ.get('SD_MOVE_DEBUG', None) is not None
 debug_move = log.trace if debug else lambda *args, **kwargs: None
-offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'omnigen2', 'cogview4', 'cosmos', 'chroma']
+offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'omnigen2', 'cogview4', 'cosmos', 'chroma', 'x-omni']
 offload_post = ['h1']
 offload_hook_instance = None
 balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline']
@@ -224,7 +224,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
                     for module_name in get_module_names(pipe):
                         module_instance = getattr(pipe, module_name, None)
                         module_cls = module_instance.__class__.__name__
-                        if (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
+                        if (module_instance is not None) and (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
                             apply_balanced_offload_to_module(module_instance, op='pre')
                 self.last_cls = module.__class__.__name__
                 process_timer.add('offload', time.time() - t0)

modules/shared_items.py

@@ -56,6 +56,7 @@ pipelines = {
     'Bria': getattr(diffusers, 'DiffusionPipeline', None),
     'hdm': getattr(diffusers, 'DiffusionPipeline', None),
     'HunyuanImage': getattr(diffusers, 'DiffusionPipeline', None),
+    'X-Omni': getattr(diffusers, 'DiffusionPipeline', None),
 }
 
 

modules/ui_extra_networks_checkpoints.py

@@ -57,7 +57,10 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
                     mtime = datetime.strptime(mtime, '%Y %B') # 2025 January
                 except Exception:
                     _size, mtime = modelstats.stat(preview_file)
-            path = f'{v.get("path", "")}+{v.get("subfolder", "")}'
+            if len(v.get("subfolder", "")) > 0:
+                path = f'{v.get("path", "")}+{v.get("subfolder", "")}'
+            else:
+                path = f'{v.get("path", "")}'
             yield {
                 "type": 'Model',
                 "name": name,

modules/video_models/video_load.py

@@ -44,24 +44,29 @@ def load_model(selected: models_def.Model):
 
     # transformer
     try:
+        def load_dit_folder(dit_folder):
+            if dit_folder is not None and dit_folder not in kwargs:
+                # get a new quant arg on every loop to prevent the quant config classes getting entangled
+                load_args, quant_args = model_quant.get_dit_args({}, module='Model', device_map=True)
+                shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
+                kwargs[dit_folder] = selected.dit_cls.from_pretrained(
+                    pretrained_model_name_or_path=selected.dit or selected.repo,
+                    subfolder=dit_folder,
+                    revision=selected.dit_revision or selected.repo_revision,
+                    cache_dir=shared.opts.hfcache_dir,
+                    **load_args,
+                    **quant_args
+                )
+            else:
+                shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} skip')
+
         if selected.dit_folder is None:
             selected.dit_folder = ['transformer']
         if isinstance(selected.dit_folder, list) or isinstance(selected.dit_folder, tuple):
             for dit_folder in selected.dit_folder: # wan a14b has transformer and transformer_2
-                if dit_folder is not None and dit_folder not in kwargs:
-                    # get a new quant arg on every loop to prevent the quant config classes getting entangled
-                    load_args, quant_args = model_quant.get_dit_args({}, module='Model', device_map=True)
-                    shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
-                    kwargs[dit_folder] = selected.dit_cls.from_pretrained(
-                        pretrained_model_name_or_path=selected.dit or selected.repo,
-                        subfolder=dit_folder,
-                        revision=selected.dit_revision or selected.repo_revision,
-                        cache_dir=shared.opts.hfcache_dir,
-                        **load_args,
-                        **quant_args
-                    )
-                else:
-                    shared.log.debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" module="{dit_folder}" folder="{dit_folder}" cls={selected.dit_cls.__name__} skip')
+                load_dit_folder(dit_folder)
+        else:
+            load_dit_folder(selected.dit_folder)
     except Exception as e:
         shared.log.error(f'video load: module=transformer cls={selected.dit_cls.__name__} {e}')
         errors.display(e, 'video')

pipelines/model_xomni.py

@@ -0,0 +1,113 @@
+import torch
+import transformers
+import diffusers
+from modules import shared, devices, sd_models, model_quant
+
+
+class XOmniPipeline(diffusers.DiffusionPipeline):
+    def __init__(
+            self,
+            tokenizer=None,
+            model=None,
+        ):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.model = model
+        self.register_modules(
+            tokenizer=tokenizer,
+            model=model,
+        )
+
+    def load(
+            self,
+            repo_id,
+            load_config: dict = {},
+        ):
+        from pipelines.xomni import modeling_xomni
+        load_args, quant_args = model_quant.get_dit_args(load_config, module='Model', device_map=True)
+        shared.log.debug(f'Load model: cls=XOmniPipeline module=tokenizer repo_id="{repo_id}"')
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(
+            repo_id,
+            use_fast=True,
+        )
+        shared.log.debug(f'Load model: cls=XOmniPipeline module=transformer repo_id="{repo_id}" args={load_args}')
+        # self.model = transformers.AutoModelForCausalLM.from_pretrained(
+        self.model = modeling_xomni.XOmniForCausalLM.from_pretrained(
+            repo_id,
+            # trust_remote_code=True,
+            cache_dir=shared.opts.hfcache_dir,
+            **load_args,
+            **quant_args,
+        )
+        flux_repo_id = "black-forest-labs/FLUX.1-dev"
+        shared.log.debug(f'Load model: cls=XOmniPipeline module=vision repo_id="{flux_repo_id}"')
+        self.model.init_vision(
+            flux_repo_id,
+            **quant_args,
+        )
+        self.model.set_generation_mode('image')
+
+    def __call__(
+            self,
+            prompt: str = "",
+            width: int = 1024,
+            height: int = 1024,
+            seed: int = -1,
+            temperature: float = 1.0,
+            downsample_size: int = 16,
+            min_p: float = 0.03,
+            top_p: float = 1.0,
+            cfg_scale: float = 1.0,
+        ):
+
+        if isinstance(prompt, list):
+            prompt = prompt[0]
+        token_h, token_w = height // downsample_size, width // downsample_size
+        image_prefix = f'<SOM>{token_h} {token_w}<IMAGE>'
+        generation_config = transformers.generation.GenerationConfig(
+            max_new_tokens=token_h * token_w,
+            do_sample=True,
+            temperature=temperature,
+            min_p=min_p,
+            top_p=top_p,
+            guidance_scale=cfg_scale,
+            suppress_tokens=self.tokenizer.convert_tokens_to_ids(self.model.config.mm_special_tokens),
+        )
+
+        # Sample inputs:
+        tokens = self.tokenizer(
+            [prompt + image_prefix],
+            return_tensors='pt',
+            padding='longest',
+            padding_side='left',
+        )
+        input_ids = tokens.input_ids.to(devices.device)
+        attention_mask = tokens.attention_mask.to(devices.device)
+        negative_ids = self.tokenizer.encode(
+            image_prefix,
+            add_special_tokens=False,
+            return_tensors='pt',
+        ).to(devices.device).expand(1, -1)
+
+        torch.manual_seed(seed)
+        tokens = self.model.generate(
+            inputs=input_ids,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            negative_prompt_ids=negative_ids,
+        )
+
+        tokens = torch.nn.functional.pad(tokens, (0, 1), value=self.tokenizer.convert_tokens_to_ids('<EOM>'))
+        torch.manual_seed(seed)
+        _, images = self.model.mmdecode(self.tokenizer, tokens[0], skip_special_tokens=False)
+        images[0].save('/tmp/xomni_out.png')
+        return images
+
+
+def load_xomni(checkpoint_info, diffusers_load_config={}):
+    repo_id = sd_models.path_to_repo(checkpoint_info)
+    sd_models.hf_auth_check(checkpoint_info)
+
+    pipe = XOmniPipeline()
+    pipe.load(repo_id, load_config=diffusers_load_config)
+    return pipe

pipelines/xomni/configuration_xomni.py

@@ -0,0 +1,25 @@
+from transformers import AutoConfig, Qwen2Config
+from typing import Tuple
+
+
+class XOmniConfig(Qwen2Config):
+    model_type = "x-omni"
+
+    def __init__(
+        self,
+        num_mm_adap_layers: int = 4,
+        num_mm_head_layers: int = 4,
+        mm_vocab_size: int = 16448,
+        image_vocab_size: int = 16384,
+        mm_special_tokens: Tuple[str] = ('<SOM>', '<EOM>', '<IMAGE>'),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_mm_adap_layers = num_mm_adap_layers
+        self.num_mm_head_layers = num_mm_head_layers
+        self.mm_vocab_size = mm_vocab_size
+        self.image_vocab_size = image_vocab_size
+        self.mm_special_tokens = mm_special_tokens
+
+
+AutoConfig.register("x-omni", XOmniConfig)

pipelines/xomni/modeling_siglip_flux.py

@@ -0,0 +1,841 @@
+import torch
+import numpy as np
+
+from typing import Any, Callable, Dict, Tuple, List, Optional, Union
+from diffusers import FluxTransformer2DModel
+from diffusers.configuration_utils import register_to_config
+from diffusers.utils import logging, USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.pipelines.flux.pipeline_flux import FluxPipeline, calculate_shift, retrieve_timesteps
+from diffusers.image_processor import PipelineImageInput
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def drop_token(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0], x.shape[1], 1)
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class FluxTransformer2DModelWithSigLIP(FluxTransformer2DModel):
+    @register_to_config
+    def __init__(
+            self,
+            patch_size: int = 1,
+            in_channels: int = 64,
+            out_channels: Optional[int] = None,
+            num_layers: int = 19,
+            num_single_layers: int = 38,
+            attention_head_dim: int = 128,
+            num_attention_heads: int = 24,
+            joint_attention_dim: int = 4096,
+            pooled_projection_dim: int = 768,
+            guidance_embeds: bool = False,
+            axes_dims_rope: Tuple[int] = (16, 56, 56),
+            siglip_channels: Optional[int] = None,
+            drop_token_prob: float = 0.,
+    ):
+        super().__init__(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            num_layers=num_layers,
+            num_single_layers=num_single_layers,
+            attention_head_dim=attention_head_dim,
+            num_attention_heads=num_attention_heads,
+            joint_attention_dim=joint_attention_dim,
+            pooled_projection_dim=pooled_projection_dim,
+            guidance_embeds=guidance_embeds,
+            axes_dims_rope=axes_dims_rope,
+        )
+        self.drop_token_prob = drop_token_prob
+        if siglip_channels is not None:
+            self.init_siglip_embed(siglip_channels)
+
+    def init_siglip_embed(self, siglip_channels):
+        self.siglip_embed = torch.nn.Linear(siglip_channels, self.inner_dim, bias=False)
+        torch.nn.init.zeros_(self.siglip_embed.weight)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            pooled_projections: torch.Tensor = None,
+            timestep: torch.LongTensor = None,
+            img_ids: torch.Tensor = None,
+            txt_ids: torch.Tensor = None,
+            guidance: torch.Tensor = None,
+            siglip_tensor: Optional[torch.Tensor] = None,
+            joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+            controlnet_block_samples=None,
+            controlnet_single_block_samples=None,
+            return_dict: bool = True,
+            controlnet_blocks_repeat: bool = False,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if txt_ids.ndim == 3:
+            logger.warning(
+                "Passing `txt_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            logger.warning(
+                "Passing `img_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            img_ids = img_ids[0]
+
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        image_rotary_emb = self.pos_embed(ids)
+
+        if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+            ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+            ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+            joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                            hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        if siglip_tensor is not None:
+            siglip_tensor = drop_token(siglip_tensor, self.drop_token_prob, training=self.training)
+            hidden_states = hidden_states + self.siglip_embed(siglip_tensor)
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1]:, ...] = (
+                        hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+                        + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+
+def teacache_forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+        siglip_tensor: Optional[torch.Tensor] = None,
+) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+    """
+    The [`FluxTransformer2DModel`] forward method.
+
+    Args:
+        hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+            Input `hidden_states`.
+        encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+            Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+        pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+            from the embeddings of input conditions.
+        timestep ( `torch.LongTensor`):
+            Used to indicate denoising step.
+        block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+            A list of tensors that if specified are added to the residuals of transformer blocks.
+        joint_attention_kwargs (`dict`, *optional*):
+            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+            `self.processor` in
+            [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+            tuple.
+
+    Returns:
+        If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+        `tuple` where the first element is the sample tensor.
+    """
+    if joint_attention_kwargs is not None:
+        joint_attention_kwargs = joint_attention_kwargs.copy()
+        lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+    else:
+        lora_scale = 1.0
+
+    if USE_PEFT_BACKEND:
+        # weight the lora layers by setting `lora_scale` for each PEFT layer
+        scale_lora_layers(self, lora_scale)
+    else:
+        if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+            logger.warning(
+                "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+            )
+
+    batch_size, seq_len, channels = hidden_states.shape
+    device, dtype = hidden_states.device, hidden_states.dtype
+    hidden_states = self.x_embedder(hidden_states)
+
+    timestep = timestep.to(hidden_states.dtype) * 1000
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+    else:
+        guidance = None
+
+    temb = (
+        self.time_text_embed(timestep, pooled_projections)
+        if guidance is None
+        else self.time_text_embed(timestep, guidance, pooled_projections)
+    )
+    encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+    if txt_ids.ndim == 3:
+        logger.warning(
+            "Passing `txt_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        txt_ids = txt_ids[0]
+    if img_ids.ndim == 3:
+        logger.warning(
+            "Passing `img_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        img_ids = img_ids[0]
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    image_rotary_emb = self.pos_embed(ids)
+
+    if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+        ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+        ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+        joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+    if self.enable_teacache:
+        inp = hidden_states.clone()
+        temb_ = temb.clone()
+        modulated_inp, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.transformer_blocks[0].norm1(inp, emb=temb_)
+        if self.cnt == 0 or self.cnt == self.num_steps - 1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else:
+            coefficients = [4.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 2.64230861e-01]
+            rescale_func = np.poly1d(coefficients)
+            # rescale_func = Polynomial(coefficients.reverse())
+            self.accumulated_rel_l1_distance += rescale_func(((modulated_inp - self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = modulated_inp
+        self.cnt += 1
+        if self.cnt == self.num_steps:
+            self.cnt = 0
+
+    if self.enable_teacache:
+        if not should_calc:
+            hidden_states += self.previous_residual
+        else:
+            ori_hidden_states = hidden_states.clone()
+            for index_block, block in enumerate(self.transformer_blocks):
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+                    encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                        block,
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        image_rotary_emb,
+                    )
+
+                else:
+                    encoder_hidden_states, hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+
+                # controlnet residual
+                if controlnet_block_samples is not None:
+                    interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                    interval_control = int(np.ceil(interval_control))
+                    # For Xlabs ControlNet.
+                    if controlnet_blocks_repeat:
+                        hidden_states = (
+                                hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                        )
+                    else:
+                        hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+            if siglip_tensor is not None:
+                siglip_tensor = drop_token(siglip_tensor, self.drop_token_prob, training=self.training)
+                hidden_states = hidden_states + self.siglip_embed(siglip_tensor)
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+            for index_block, block in enumerate(self.single_transformer_blocks):
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+                    hidden_states = self._gradient_checkpointing_func(
+                        block,
+                        hidden_states,
+                        temb,
+                        image_rotary_emb,
+                    )
+
+                else:
+                    hidden_states = block(
+                        hidden_states=hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+
+                # controlnet residual
+                if controlnet_single_block_samples is not None:
+                    interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                    interval_control = int(np.ceil(interval_control))
+                    hidden_states[:, encoder_hidden_states.shape[1]:, ...] = (
+                            hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+                            + controlnet_single_block_samples[index_block // interval_control]
+                    )
+
+            hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+            self.previous_residual = hidden_states - ori_hidden_states
+    else:
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                            hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+        if siglip_tensor is not None:
+            siglip_tensor = drop_token(siglip_tensor, self.drop_token_prob, training=self.training)
+            hidden_states = hidden_states + self.siglip_embed(siglip_tensor)
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1]:, ...] = (
+                        hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+                        + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+
+    hidden_states = self.norm_out(hidden_states, temb)
+    output = self.proj_out(hidden_states)
+
+    if USE_PEFT_BACKEND:
+        # remove `lora_scale` from each PEFT layer
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (output,)
+
+    return Transformer2DModelOutput(sample=output)
+
+
+class FluxPipelineWithSigLIP(FluxPipeline):
+   
+    @torch.no_grad()
+    def __call__(
+            self,
+            siglip_tensor: torch.Tensor,
+            prompt: Union[str, List[str]] = None,
+            prompt_2: Optional[Union[str, List[str]]] = None,
+            negative_prompt: Union[str, List[str]] = None,
+            negative_prompt_2: Optional[Union[str, List[str]]] = None,
+            true_cfg_scale: float = 1.0,
+            true_cfg_scale_2: float = 1.0,
+            height: Optional[int] = None,
+            width: Optional[int] = None,
+            num_inference_steps: int = 28,
+            sigmas: Optional[List[float]] = None,
+            guidance_scale: float = 3.5,
+            num_images_per_prompt: Optional[int] = 1,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.FloatTensor] = None,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            ip_adapter_image: Optional[PipelineImageInput] = None,
+            ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+            negative_ip_adapter_image: Optional[PipelineImageInput] = None,
+            negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+            negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            output_type: Optional[str] = "pil",
+            return_dict: bool = True,
+            joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+            callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+            callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+            max_sequence_length: int = 512,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
+                not greater than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            true_cfg_scale (`float`, *optional*, defaults to 1.0):
+                When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 3.5):
+                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.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_ip_adapter_image:
+                (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+        assert true_cfg_scale == true_cfg_scale_2
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        has_neg_prompt = negative_prompt is not None or (
+                negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        assert do_true_cfg
+        (
+            negative_prompt_embeds,
+            negative_pooled_prompt_embeds,
+            _,
+        ) = self.encode_prompt(
+            prompt=negative_prompt,
+            prompt_2=negative_prompt_2,
+            prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0] * 2)
+        else:
+            guidance = None
+
+        if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
+                negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
+        ):
+            negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+            negative_ip_adapter_image = [negative_ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
+
+        elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
+                negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
+        ):
+            ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+            ip_adapter_image = [ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        image_embeds = None
+        negative_image_embeds = None
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+        if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
+            negative_image_embeds = self.prepare_ip_adapter_image_embeds(
+                negative_ip_adapter_image,
+                negative_ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+                if image_embeds is not None:
+                    self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0] * 2).to(latents.dtype)
+
+                batch_noise_pred = self.transformer(
+                    hidden_states=torch.cat([latents, latents], dim=0),
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=torch.cat([pooled_prompt_embeds, negative_pooled_prompt_embeds.expand_as(pooled_prompt_embeds)], dim=0),
+                    encoder_hidden_states=torch.cat([prompt_embeds, negative_prompt_embeds.expand_as(prompt_embeds)], dim=0),
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    siglip_tensor=torch.cat([siglip_tensor, torch.zeros_like(siglip_tensor)], dim=0),
+                    return_dict=False,
+                )[0]
+                noise_pred, neg_noise_pred = batch_noise_pred.chunk(2)
+                noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)

pipelines/xomni/modeling_siglip_tokenizer.py

@@ -0,0 +1,231 @@
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import einsum
+from torchvision import transforms
+
+from PIL import Image
+from einops import rearrange
+
+from .modeling_vit import create_siglip_vit
+
+
+def create_anyres_preprocess(
+    short_size=384, 
+    long_size=1152, 
+    patch_size=16, 
+    random_ratio=None, 
+    min_short_size=128, 
+    max_aspect_ratio=3., 
+    filtering=True
+):
+
+    def resize_and_filtering(pil_image):
+        pil_image = pil_image.convert('RGB')
+        width, height = pil_image.size
+        ss, ls = min(width, height), max(width, height)
+        aspect_ratio = ls / ss
+        if filtering and (ss < min_short_size or aspect_ratio > max_aspect_ratio):
+            return None
+        target_width, target_height = width, height
+        if random_ratio is not None:
+            log_ratio = torch.log(torch.tensor(random_ratio))
+            sqrt_ratio = torch.exp(0.5 * torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()
+            target_width = int(round(target_width * sqrt_ratio))
+            target_height = int(round(target_height / sqrt_ratio))
+        
+        ss = min(target_width, target_height)
+        if ss < short_size:
+            target_width = target_width * (short_size / ss)
+            target_height = target_height * (short_size / ss)
+        
+        ls = max(target_width, target_height)
+        if ls > long_size:
+            target_width = target_width * (long_size / ls)
+            target_height = target_height * (long_size / ls)
+        
+        target_width = int(round(target_width / patch_size)) * patch_size
+        target_height = int(round(target_height / patch_size)) * patch_size
+        pil_image = pil_image.resize((target_width, target_height), resample=Image.BICUBIC)
+        
+        to_tensor = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ])
+        return to_tensor(pil_image)
+
+    transform = transforms.Lambda(resize_and_filtering)
+    return transform
+
+
+class IBQ(nn.Module):
+    def __init__(self, n_e, e_dim, skip_quantization_prob=0.0, quantization_temp=2.0, beta=0.25, sane_index_shape=False, l2_norm=True):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.quantization_temp = quantization_temp
+        self.skip_quantization_prob = skip_quantization_prob
+        self.beta = beta
+        self.sane_index_shape = sane_index_shape
+        self.l2_norm = l2_norm
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+        if self.l2_norm:
+            self.embedding.weight.data = F.normalize(self.embedding.weight.data, p=2, dim=-1)
+    
+    def forward(self, z, temp=None, rescale_logits=False, return_logits=False, **kwargs):
+        assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel"
+        assert rescale_logits == False, "Only for interface compatible with Gumbel"
+        assert return_logits == False, "Only for interface compatible with Gumbel"
+        # reshape z -> (batch, height, width, channel) and flatten
+        z = rearrange(z, 'b c h w -> b h w c').contiguous()
+        assert z.shape[-1] == self.e_dim
+        z_flattened = z.view(-1, self.e_dim)
+        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+
+        if self.l2_norm:
+            z = F.normalize(z, p=2, dim=-1)
+            z_flattened = F.normalize(z_flattened, p=2, dim=-1)
+            embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+        else:
+            embedding = self.embedding.weight
+
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(embedding**2, dim=1) - 2 * \
+            torch.einsum('bd,dn->bn', z_flattened, torch.einsum('n d -> d n', embedding))
+        
+        if self.training:
+            logits = -d / self.quantization_temp
+            soft_one_hot = F.softmax(logits, dim=1)
+            min_encoding_indices = soft_one_hot.max(1, keepdim=True)[1]
+            hard_one_hot = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(1, min_encoding_indices, 1.0)
+            one_hot = hard_one_hot - soft_one_hot.detach() + soft_one_hot
+
+            z_q = einsum('b n, n d -> b d', one_hot, self.embedding.weight).view(z.shape)
+            z_q_2 = einsum('b n, n d -> b d', hard_one_hot, self.embedding.weight).view(z.shape)
+
+            # compute loss for embedding
+            commit_loss = torch.mean((z_q - z) ** 2) + torch.mean((z_q_2.detach() - z) ** 2) + self.beta * \
+                        torch.mean((z_q_2 - z.detach()) ** 2)
+        else:
+            min_encoding_indices = torch.argmin(d, dim=1)
+            z_q = embedding[min_encoding_indices].view(z.shape)
+            commit_loss = None
+        
+        if self.training and self.skip_quantization_prob > 0.0:
+            z_q = torch.where(
+                torch.rand_like(z_q[:, 0:1, 0:1, 0:1]).expand_as(z_q) <= self.skip_quantization_prob,
+                z, z_q,
+            )
+        
+        # reshape back to match original input shape
+        z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()
+
+        if self.sane_index_shape:
+            min_encoding_indices = min_encoding_indices.reshape(z_q.shape[0], z_q.shape[2], z_q.shape[3])
+
+        return (z_q, None, min_encoding_indices), commit_loss
+
+    def get_codebook_entry(self, indices, bhwc):
+        # shape specifying (batch, height, width, channel)
+        # get quantized latent vectors
+        z_q = self.embedding(indices)
+
+        if bhwc is not None:
+            z_q = z_q.view(bhwc)
+            # reshape back to match original input shape
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels, num_groups=32):
+        super().__init__()
+        self.conv1 = nn.Conv2d(channels, channels, 3, padding='same')
+        self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=channels)
+        self.activate = nn.GELU()
+        self.conv2 = nn.Conv2d(channels, channels, 3, padding='same')
+        self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=channels)
+    
+    def forward(self, x):
+        res = x
+        x = self.norm1(x)
+        x = self.activate(x)
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.activate(x)
+        x = self.conv2(x)
+        return x + res
+
+
+class VQConvProjector(nn.Module):
+    def __init__(
+        self, 
+        z_channels=1536, 
+        codebook_size=16384, 
+        codebook_dim=2048, 
+        conv_layers=2,
+        with_norm=True,
+        skip_quant_prob=0.1,
+    ):
+        super().__init__()
+        self.quant_conv = nn.Conv2d(z_channels, codebook_dim, 1)
+        self.quantize = IBQ(codebook_size, codebook_dim, skip_quant_prob, sane_index_shape=True)
+        self.post_quant_conv = nn.Conv2d(codebook_dim, z_channels, 1)
+        block = ResidualBlock
+        self.post_conv = nn.Sequential(*[block(z_channels) for _ in range(conv_layers)])
+    
+    def forward(self, x, h, w):
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
+        z = self.quant_conv(x)
+        (z_q, _, _), codebook_loss = self.quantize(z)
+        z = self.post_quant_conv(z_q)
+        z = self.post_conv(z)
+        z = rearrange(z, 'b c h w -> b (h w) c')
+        return z, codebook_loss
+    
+    def encode(self, x, h, w):
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
+        z = self.quant_conv(x)
+        (_, _, tokens), _ = self.quantize(z)
+        return tokens
+    
+    def decode(self, tokens, bhwc):
+        z_q = self.quantize.get_codebook_entry(tokens, bhwc)
+        z = self.post_quant_conv(z_q)
+        z = self.post_conv(z)        
+        return z
+
+
+class SiglipTokenizer(nn.Module):
+    def __init__(
+        self, 
+        siglip_name, 
+        siglip_path, 
+        projector_path, 
+        z_channels=1536, 
+        codebook_size=16384, 
+        codebook_dim=2048, 
+        with_norm=True
+    ):
+        super().__init__()
+        self.vit = create_siglip_vit(model_name=siglip_name, path=siglip_path)
+        self.vqproj = VQConvProjector(
+            z_channels=z_channels, 
+            codebook_size=codebook_size, 
+            codebook_dim=codebook_dim, 
+            with_norm=with_norm
+        )
+        self.vqproj.load_state_dict(torch.load(projector_path, map_location='cpu'), strict=True)
+
+    def encode(self, x):
+        features, (h, w), _ = self.vit(x)
+        tokens = self.vqproj.encode(features, h, w)
+        return tokens
+    
+    def decode(self, tokens, bhwc):
+        return self.vqproj.decode(tokens, bhwc)

pipelines/xomni/modeling_vit.py

@@ -0,0 +1,699 @@
+import math
+import warnings
+from dataclasses import dataclass
+from functools import partial
+from typing import (
+    Callable, Dict, Final, List, Literal, Optional,
+    Sequence, Set, Tuple, Type, Union,
+)
+
+from torch.utils.checkpoint import checkpoint
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.layers import (
+    DropPath, LayerType, Mlp, PatchDropout,
+    PatchEmbed, resample_abs_pos_embed,
+)
+from timm.models._manipulate import checkpoint_seq, named_apply
+
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+            "The distribution of values may be incorrect.",
+            stacklevel=2,
+        )
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)  # noqa: E741
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.0))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    # type: (torch.Tensor, float, float, float, float) -> torch.Tensor
+    r"""The original timm.models.layers.weight_init.trunc_normal_ can not handle bfloat16 yet, here we first
+    convert the tensor to float32, apply the trunc_normal_() in float32, and then convert it back to its orignal dtype.
+    Fills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn
+    from the normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+
+    with torch.no_grad():
+        dtype = tensor.dtype
+        tensor_fp32 = tensor.float()
+        tensor_fp32 = _no_grad_trunc_normal_(tensor_fp32, mean, std, a, b)
+        tensor_dtype = tensor_fp32.to(dtype=dtype)
+        tensor.copy_(tensor_dtype)
+
+
+def init_weights(self):
+    if self.pos_embed is not None:
+        trunc_normal_(self.pos_embed, std=self.pos_embed.shape[1] ** -0.5)
+    trunc_normal_(self.latent, std=self.latent_dim**-0.5)
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = "") -> None:
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+    elif hasattr(module, "init_weights"):
+        module.init_weights()
+
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim**-0.5
+        # self.fused_attn = use_fused_attn()
+        self.fused_attn = True
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop) if proj_drop > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, cu_slens=None) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if cu_slens is not None:
+            q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
+            k = k.permute(0, 2, 1, 3)
+            v = v.permute(0, 2, 1, 3)
+            max_seqlen = torch.max(cu_slens[1:] - cu_slens[:-1]).item()
+            x = flash_attn_varlen_func(
+                q.squeeze(0),
+                k.squeeze(0),
+                v.squeeze(0),
+                cu_seqlens_q=cu_slens,
+                cu_seqlens_k=cu_slens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                softmax_scale=self.scale,
+                causal=False,
+                )
+
+            x = x.reshape(B, N, -1)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+
+        else:
+            q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
+            k = k.permute(0, 2, 1, 3)
+            v = v.permute(0, 2, 1, 3)
+            x = flash_attn_func(q, k, v, softmax_scale=self.scale) # -> b, n, h, c
+
+            x = x.reshape(B, N, -1)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: float = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        proj_drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values: Optional[float] = None,
+        drop_path: float = 0.0,
+        act_layer: nn.Module = nn.GELU,
+        norm_layer: nn.Module = nn.LayerNorm,
+        mlp_layer: nn.Module = Mlp,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+        )
+        self.ls1 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = mlp_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=proj_drop,
+        )
+        self.ls2 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor, cu_slens=None) -> torch.Tensor:
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x), cu_slens=cu_slens)))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer
+
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
+        - https://arxiv.org/abs/2010.11929
+    """
+
+    dynamic_img_size: Final[bool]
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        num_classes: int = 1000,
+        global_pool: Literal["", "avg", "token", "map"] = "token",
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        init_values: Optional[float] = None,
+        class_token: bool = True,
+        no_embed_class: bool = False,
+        reg_tokens: int = 0,
+        pre_norm: bool = False,
+        fc_norm: Optional[bool] = None,
+        dynamic_img_size: bool = False,
+        dynamic_img_pad: bool = False,
+        drop_rate: float = 0.0,
+        pos_drop_rate: float = 0.0,
+        patch_drop_rate: float = 0.0,
+        proj_drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        weight_init: Literal["skip", "jax", "jax_nlhb", "moco", ""] = "",
+        embed_layer: Callable = PatchEmbed,
+        norm_layer: Optional[LayerType] = None,
+        act_layer: Optional[LayerType] = None,
+        strict_img_size: bool = False,
+        block_fn: Type[nn.Module] = Block,
+        mlp_layer: Type[nn.Module] = Mlp,
+        ignore_head: bool = False,
+    ) -> None:
+        """
+        Args:
+            img_size: Input image size.
+            patch_size: Patch size.
+            in_chans: Number of image input channels.
+            num_classes: Mumber of classes for classification head.
+            global_pool: Type of global pooling for final sequence (default: 'token').
+            embed_dim: Transformer embedding dimension.
+            depth: Depth of transformer.
+            num_heads: Number of attention heads.
+            mlp_ratio: Ratio of mlp hidden dim to embedding dim.
+            qkv_bias: Enable bias for qkv projections if True.
+            init_values: Layer-scale init values (layer-scale enabled if not None).
+            class_token: Use class token.
+            no_embed_class: Don't include position embeddings for class (or reg) tokens.
+            reg_tokens: Number of register tokens.
+            fc_norm: Pre head norm after pool (instead of before), if None, enabled when global_pool == 'avg'.
+            drop_rate: Head dropout rate.
+            pos_drop_rate: Position embedding dropout rate.
+            attn_drop_rate: Attention dropout rate.
+            drop_path_rate: Stochastic depth rate.
+            weight_init: Weight initialization scheme.
+            embed_layer: Patch embedding layer.
+            norm_layer: Normalization layer.
+            act_layer: MLP activation layer.
+            block_fn: Transformer block layer.
+        """
+        super().__init__()
+        assert global_pool in ("", "avg", "token", "map")
+        assert class_token or global_pool != "token"
+        use_fc_norm = global_pool == "avg" if fc_norm is None else fc_norm
+        # norm_layer = get_norm_layer(norm_layer) or partial(nn.LayerNorm, eps=1e-6)
+        # act_layer = get_act_layer(act_layer) or nn.GELU
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        act_layer = nn.GELU
+
+        self.num_classes = num_classes
+        self.global_pool = global_pool
+        self.num_features = self.embed_dim = (
+            embed_dim  # num_features for consistency with other models
+        )
+        self.num_prefix_tokens = 1 if class_token else 0
+        self.num_prefix_tokens += reg_tokens
+        self.num_reg_tokens = reg_tokens
+        self.has_class_token = class_token
+        self.no_embed_class = (
+            no_embed_class  # don't embed prefix positions (includes reg)
+        )
+        self.dynamic_img_size = dynamic_img_size
+        self.grad_checkpointing = False
+        self.ignore_head = ignore_head
+
+        embed_args = {}
+        if dynamic_img_size:
+            # flatten deferred until after pos embed
+            embed_args.update(dict(strict_img_size=False, output_fmt="NHWC"))
+        self.patch_embed = embed_layer(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            bias=not pre_norm,  # disable bias if pre-norm is used (e.g. CLIP)
+            dynamic_img_pad=dynamic_img_pad,
+            strict_img_size=strict_img_size,
+            **embed_args,
+        )
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = (
+            nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None
+        )
+        self.reg_token = (
+            nn.Parameter(torch.zeros(1, reg_tokens, embed_dim)) if reg_tokens else None
+        )
+        embed_len = (
+            num_patches if no_embed_class else num_patches + self.num_prefix_tokens
+        )
+        self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * 0.02)
+        self.pos_drop = nn.Dropout(p=pos_drop_rate)
+        if patch_drop_rate > 0:
+            self.patch_drop = PatchDropout(
+                patch_drop_rate,
+                num_prefix_tokens=self.num_prefix_tokens,
+            )
+        else:
+            self.patch_drop = nn.Identity()
+        self.norm_pre = norm_layer(embed_dim) if pre_norm else nn.Identity()
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(
+            *[
+                block_fn(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_norm=qk_norm,
+                    init_values=init_values,
+                    proj_drop=proj_drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    act_layer=act_layer,
+                    mlp_layer=mlp_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+
+    def init_weights(self, mode: Literal["jax", "jax_nlhb", "moco", ""] = "") -> None:
+        assert mode in ("jax", "jax_nlhb", "moco", "")
+        # head_bias = -math.log(self.num_classes) if "nlhb" in mode else 0.0
+        trunc_normal_(self.pos_embed, std=0.02)
+        if self.cls_token is not None:
+            nn.init.normal_(self.cls_token, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    @torch.jit.ignore
+    def no_weight_decay(self) -> Set:
+        return {"pos_embed", "cls_token", "dist_token"}
+
+    @torch.jit.ignore
+    def group_matcher(self, coarse: bool = False) -> Dict:
+        return dict(
+            stem=r"^cls_token|pos_embed|patch_embed",  # stem and embed
+            blocks=[(r"^blocks\.(\d+)", None), (r"^norm", (99999,))],
+        )
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable: bool = True) -> None:
+        self.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def get_classifier(self) -> nn.Module:
+        return self.head
+
+    def reset_classifier(self, num_classes: int, global_pool=None) -> None:
+        self.num_classes = num_classes
+        if global_pool is not None:
+            assert global_pool in ("", "avg", "token", "map")
+            if global_pool == "map" and self.attn_pool is None:
+                assert (
+                    False
+                ), "Cannot currently add attention pooling in reset_classifier()."
+            elif global_pool != "map " and self.attn_pool is not None:
+                self.attn_pool = None  # remove attention pooling
+            self.global_pool = global_pool
+        self.head = (
+            nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        )
+
+    def rescale_positional_embedding(self, out_size):
+        h, w = out_size
+        pos_embed_shape = int((self.pos_embed.shape[1]) ** 0.5)
+        if (h, w) == (pos_embed_shape, pos_embed_shape):
+            return self.pos_embed
+        rescaled_positional_embedding = \
+            self.pos_embed.new_zeros(1, h*w, self.pos_embed.shape[2])
+        pe_2d = self.pos_embed[0].T.contiguous().view(1, -1, pos_embed_shape, pos_embed_shape)
+        pe_2d = F.interpolate(pe_2d, out_size, mode='bilinear', align_corners=False).view(-1, h*w)
+        rescaled_positional_embedding[0] = pe_2d.T.contiguous()
+        return rescaled_positional_embedding
+
+    def _pos_embed(self, x: torch.Tensor) -> torch.Tensor:
+        if self.dynamic_img_size:
+            B, H, W, C = x.shape
+            pos_embed = resample_abs_pos_embed(
+                self.pos_embed,
+                (H, W),
+                num_prefix_tokens=0 if self.no_embed_class else self.num_prefix_tokens,
+            )
+            x = x.view(B, -1, C)
+        else:
+            pos_embed = self.pos_embed
+
+        to_cat = []
+        if self.cls_token is not None:
+            to_cat.append(self.cls_token.expand(x.shape[0], -1, -1))
+        if self.reg_token is not None:
+            to_cat.append(self.reg_token.expand(x.shape[0], -1, -1))
+
+        if self.no_embed_class:
+            # deit-3, updated JAX (big vision)
+            # position embedding does not overlap with class token, add then concat
+            x = x + pos_embed
+            if to_cat:
+                x = torch.cat(to_cat + [x], dim=1)
+        else:
+            # original timm, JAX, and deit vit impl
+            # pos_embed has entry for class token, concat then add
+            if to_cat:
+                x = torch.cat(to_cat + [x], dim=1)
+            x = x + pos_embed
+
+        return self.pos_drop(x)
+
+    def _intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,
+    ) -> List[torch.Tensor]:
+        outputs, num_blocks = [], len(self.blocks)
+        take_indices = set(
+            range(num_blocks - n, num_blocks) if isinstance(n, int) else n
+        )
+
+        # forward pass
+        x = self.patch_embed(x)
+        x = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in take_indices:
+                outputs.append(x)
+
+        return outputs
+
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, Sequence] = 1,
+        reshape: bool = False,
+        return_prefix_tokens: bool = False,
+        norm: bool = False,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        """Intermediate layer accessor (NOTE: This is a WIP experiment).
+        Inspired by DINO / DINOv2 interface
+        """
+        # take last n blocks if n is an int, if in is a sequence, select by matching indices
+        outputs = self._intermediate_layers(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        prefix_tokens = [out[:, 0 : self.num_prefix_tokens] for out in outputs]
+        outputs = [out[:, self.num_prefix_tokens :] for out in outputs]
+
+        if reshape:
+            grid_size = self.patch_embed.grid_size
+            outputs = [
+                out.reshape(x.shape[0], grid_size[0], grid_size[1], -1)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+                for out in outputs
+            ]
+
+        if return_prefix_tokens:
+            return tuple(zip(outputs, prefix_tokens))
+        return tuple(outputs)
+
+    def forward_features_list(self, x_list):
+        x_all = []
+        image_sizes = []
+        for x in x_list:
+            bs, _, h, w = x.shape
+
+            # fix patch size=14 in datasets 
+            pad_h = (self.patch_embed.patch_size[0] - h % self.patch_embed.patch_size[0]) % self.patch_embed.patch_size[0]
+            pad_w = (self.patch_embed.patch_size[1] - w % self.patch_embed.patch_size[1]) % self.patch_embed.patch_size[1]
+            x = F.pad(x, (0, pad_w, 0, pad_h))
+
+            bs, _, h, w = x.shape
+
+            h = h // self.patch_embed.patch_size[0]
+            w = w // self.patch_embed.patch_size[1]
+
+            x = self.patch_embed(x)
+            x = x + self.rescale_positional_embedding(out_size=(h, w))
+            x = self.patch_drop(x)
+            x = self.norm_pre(x)
+            x_all.append(x)
+            image_sizes.append((h, w))
+
+        slen = [xi.size(1) for xi in x_all]
+        x = torch.cat(x_all, dim=1)
+
+        cu_indices = [0, ]
+        for i in slen:
+            cu_indices.append(cu_indices[-1] + i)
+
+        cu_slens = torch.tensor(cu_indices, dtype=torch.int32).to(x.device)
+        for idx, blk in enumerate(self.blocks):
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(blk, x, cu_slens, use_reentrant=True)
+            else:
+                x = blk(x, cu_slens=cu_slens)
+        feats = x.split(slen, dim=1) #[(1, slen, c)]
+        return feats, image_sizes
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        bs, _, h, w = x.shape
+        h = h // self.patch_embed.patch_size[0]
+        w = w // self.patch_embed.patch_size[1]
+        
+        x = self.patch_embed(x)
+        # x = self._pos_embed(x)
+        x = x + self.rescale_positional_embedding(out_size=(h, w))
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+        if self.grad_checkpointing and not torch.jit.is_scripting():
+            x = checkpoint_seq(self.blocks, x)
+        else:
+            x = self.blocks(x)
+        return x, (h, w)
+
+    def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor:
+        x = self.norm(x)
+        if self.attn_pool is not None:
+            x = self.attn_pool(x)
+        elif self.global_pool == "avg":
+            x = x[:, self.num_prefix_tokens :].mean(dim=1)
+        elif self.global_pool:
+            x = x[:, 0]  # class token
+        x = self.fc_norm(x)
+        x = self.head_drop(x)
+        return x if pre_logits else self.head(x)
+
+    def forward(self, x, cal_attn_pool=False):
+        if type(x) is list:
+            x, image_sizes = self.forward_features_list(x)
+            return x, image_sizes, None
+        else:
+            x, image_sizes = self.forward_features(x)
+            return x, image_sizes, None
+
+@dataclass
+class SigLIPVisionCfg:
+    width: int = 1152
+    layers: Union[Tuple[int, int, int, int], int] = 27
+    heads: int = 16
+    patch_size: int = 14
+    image_size: Union[Tuple[int, int], int] = 336
+    global_pool: str = "map"
+    mlp_ratio: float = 3.7362
+    class_token: bool = False
+    num_classes: int = 0
+    use_checkpoint: bool = False
+
+
+SigLIP_MODEL_CONFIG = {
+    "siglip_so400m_patch16_384": {
+        "image_size": 384,
+        "patch_size": 16,
+        "width": 1152,
+        "layers": 27,
+        "heads": 16,
+        "mlp_ratio": 3.7362,
+        "global_pool": "map",
+        "use_checkpoint": False,
+    },
+    "siglip2_giant_patch16_384":{
+        "image_size": 384,
+        "patch_size": 16,
+        "width": 1536,
+        "layers": 40,
+        "heads": 16,
+        "mlp_ratio": 4,
+        "global_pool": "map",
+        "use_checkpoint": False,
+    },
+}
+
+
+def resize_evaclip_pos_embed(model: VisionTransformer, interpolation: str = 'bicubic'):
+    # interpolate position embedding
+    orig_size = 24
+    new_size = 128
+    pos_tokens = model.pos_embed
+    pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, model.embed_dim).permute(0, 3, 1, 2)
+    pos_tokens = torch.nn.functional.interpolate(
+        pos_tokens, size=(new_size, new_size), mode=interpolation, align_corners=False)
+    pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+    model.pos_embed = nn.Parameter(pos_tokens, requires_grad=True)
+    return model 
+
+
+def create_siglip_vit(
+    model_name: str = "siglip_so400m_patch14_384",
+    select_layer: int = -1,
+    path: str = "",
+    gradient_checkpointing: bool = False,
+    **kwargs,
+):
+    vision_cfg = SigLIPVisionCfg(**SigLIP_MODEL_CONFIG[model_name])
+
+    if select_layer <= 0:
+        layers = min(vision_cfg.layers, vision_cfg.layers + select_layer + 1)
+    else:
+        layers = min(vision_cfg.layers, select_layer)
+
+    model = VisionTransformer(
+        img_size=2048,
+        patch_size=16,
+        embed_dim=vision_cfg.width,
+        depth=layers,
+        num_heads=vision_cfg.heads,
+        mlp_ratio=vision_cfg.mlp_ratio,
+        class_token=vision_cfg.class_token,
+        global_pool=vision_cfg.global_pool,
+        dynamic_img_pad=False,
+        strict_img_size=False,
+        ignore_head=kwargs.get("ignore_head", False),
+        weight_init=kwargs.get("weight_init", "skip"),
+        num_classes=0
+    )
+    model.config = vision_cfg
+    state_dict = torch.load(path, map_location="cpu")
+    model.load_state_dict(state_dict, strict=False)
+
+    if gradient_checkpointing:
+        model.set_grad_checkpointing(True)
+    return model

pipelines/xomni/modeling_xomni.py

@@ -0,0 +1,317 @@
+import os
+from types import SimpleNamespace
+from typing import Tuple, List, Optional, Union
+
+import torch
+import torch.nn as nn
+
+from huggingface_hub import hf_hub_download
+from transformers import Qwen2ForCausalLM, AutoModel, AutoModelForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm, Qwen2RotaryEmbedding, Qwen2DecoderLayer, Qwen2Model, Qwen2PreTrainedModel
+
+from .configuration_xomni import XOmniConfig
+from .modeling_siglip_tokenizer import create_anyres_preprocess, SiglipTokenizer
+from .modeling_siglip_flux import FluxTransformer2DModelWithSigLIP, FluxPipelineWithSigLIP
+from .modeling_vit import create_siglip_vit
+
+
+class XOmniDecoderLayer(Qwen2DecoderLayer):
+    def __init__(self, config: XOmniConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.layer_idx = layer_idx
+        self.is_lm_layer = config.num_mm_adap_layers <= layer_idx < config.num_hidden_layers - config.num_mm_head_layers
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        hidden_states, multimodal_mask = torch.split(hidden_states, hidden_states.shape[-1] // 2, dim=-1)
+        if self.is_lm_layer:
+            output_hidden_states, *others = super().forward(hidden_states, **kwargs)
+            output_hidden_states = torch.cat([output_hidden_states, multimodal_mask], dim=-1)
+            return output_hidden_states, *others
+        
+        # mm_hidden_states = torch.where(multimodal_mask.bool(), hidden_states, torch.zeros_like(hidden_states))
+        output_hidden_states, *others = super().forward(hidden_states, **kwargs)
+        output_hidden_states = torch.where(multimodal_mask.bool(), output_hidden_states, hidden_states)
+        output_hidden_states = torch.cat([output_hidden_states, multimodal_mask], dim=-1)
+        return output_hidden_states, *others
+
+
+class XOmniModel(Qwen2Model, Qwen2PreTrainedModel):
+    model_type = "x-omni"
+    config_class = XOmniConfig
+
+    def __init__(self, config: XOmniConfig):
+        Qwen2PreTrainedModel.__init__(self, config)
+        self.padding_idx = -1
+        self.vocab_size = config.vocab_size
+        
+        self.lm_embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.mm_embed_tokens = nn.Embedding(config.mm_vocab_size, config.hidden_size, self.padding_idx)
+
+        self.layers = nn.ModuleList(
+            [XOmniDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self._attn_implementation = config._attn_implementation
+        self.lm_norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mm_norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen2RotaryEmbedding(config=config)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.lm_embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.lm_embed_tokens = value
+
+    def embed_tokens(self, input_ids):
+        (B, L), C = input_ids.shape, self.config.hidden_size
+        multimodal_mask = input_ids >= self.config.vocab_size
+        lm_input_ids = input_ids[~multimodal_mask][None, :]
+        mm_input_ids = input_ids[multimodal_mask][None, :] - self.config.vocab_size
+        lm_embeds = self.lm_embed_tokens(lm_input_ids)
+        mm_embeds = self.mm_embed_tokens(mm_input_ids)
+
+        inputs_embeds = lm_embeds.new_empty((B, L, C))
+        multimodal_mask = multimodal_mask[:, :, None].expand_as(inputs_embeds)
+        inputs_embeds[~multimodal_mask] = lm_embeds.reshape(-1)
+        inputs_embeds[multimodal_mask] = mm_embeds.reshape(-1)
+
+        inputs_embeds = torch.cat([inputs_embeds, multimodal_mask.to(inputs_embeds.dtype)], dim=-1)
+        return inputs_embeds
+
+    def norm(self, hidden_states):
+        hidden_states, multimodal_mask = torch.split(hidden_states, hidden_states.shape[-1] // 2, dim=-1)
+        return torch.where(multimodal_mask.bool(), self.mm_norm(hidden_states), self.lm_norm(hidden_states))
+
+
+class XOmniForCausalLM(Qwen2ForCausalLM):
+    model_type = "x-omni"
+    config_class = XOmniConfig
+    
+    _keys_to_ignore_on_load_missing = r'image_tokenizer\.*'
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = XOmniModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.mm_head = nn.Linear(config.hidden_size, config.mm_vocab_size, bias=False)
+        
+        self.generation_mode = 'text'
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @property
+    def device(self):
+        return next(iter(self.parameters())).device
+    
+    def init_vision(self, flux_pipe_path, **kwargs):
+        self.som_token = self.config.mm_special_tokens[0]
+        self.eom_token = self.config.mm_special_tokens[1]
+        self.img_token = self.config.mm_special_tokens[2]
+
+        self.vision_config = SimpleNamespace(**self.config.vision_config)
+        self.transform_config = SimpleNamespace(**self.vision_config.transform)
+        self.encoder_config = SimpleNamespace(**self.vision_config.encoder)
+        self.decoder_config = SimpleNamespace(**self.vision_config.decoder)
+
+        dtype_map = {'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16}
+        self.vision_dtype = dtype_map[self.vision_config.dtype]
+
+        self.image_transform = create_anyres_preprocess(**self.vision_config.transform)
+        
+        self.encoder_config.siglip_path = os.path.join(self.name_or_path, self.encoder_config.siglip_path) if os.path.isdir(self.name_or_path) else hf_hub_download(repo_id=self.name_or_path, filename=self.encoder_config.siglip_path)
+        self.encoder_config.projector_path = os.path.join(self.name_or_path, self.encoder_config.projector_path) if os.path.isdir(self.name_or_path) else hf_hub_download(repo_id=self.name_or_path, filename=self.encoder_config.projector_path)
+
+        self.image_tokenizer = SiglipTokenizer(**vars(self.encoder_config))
+        self.image_tokenizer.to(self.device, self.vision_dtype)
+
+        transformer = FluxTransformer2DModelWithSigLIP.from_pretrained(
+            self.name_or_path, 
+            siglip_channels=self.encoder_config.z_channels, 
+            torch_dtype=self.vision_dtype,
+            subfolder=self.decoder_config.model_path,
+            **kwargs,
+        )
+
+        self.decoder_pipe = FluxPipelineWithSigLIP.from_pretrained(
+            flux_pipe_path, 
+            transformer=transformer,
+            torch_dtype=self.vision_dtype,
+        )
+        self.decoder_pipe.set_progress_bar_config(disable=True)
+
+    def set_generation_mode(self, mode):
+        assert mode in ('text', 'image'), f'Invalid generation mode: {mode}'
+        self.generation_mode = mode
+
+    def mmencode(self, tokenizer, texts=None, images=None, **kwargs):
+        texts = texts or []
+        images = images or []
+        doc = ''
+        while len(texts) > 0 or len(images) > 0:
+            if len(texts) > 0:
+                doc += texts.pop(0)
+            if len(images) > 0:
+                doc += self.tokenize_image(images.pop(0))
+        return tokenizer.encode(doc, **kwargs)
+    
+    def mmdecode(self, tokenizer, token_ids, force_text=None, **kwargs):
+        force_text = force_text or []
+        if isinstance(token_ids, torch.Tensor):
+            if len(token_ids.shape) == 2:
+                assert token_ids.shape[0] == 1
+                token_ids = token_ids[0]
+            assert len(token_ids.shape) == 1
+        else:
+            if not isinstance(token_ids[0], int):
+                assert len(token_ids) == 1
+                token_ids = token_ids[0]
+            assert isinstance(token_ids[0], int)
+        
+        doc = tokenizer.decode(token_ids, **kwargs)
+        doc = doc.replace(tokenizer.pad_token, '')
+        doc = doc.replace('<SEP>', '')
+        texts, images = [], []
+        text_image_chunks = doc.split(self.eom_token)
+        for chunk in text_image_chunks:
+            text, image_str = chunk.split(self.som_token) \
+                if self.som_token in chunk else (chunk, '')
+            texts.append(text)
+            if self.img_token in image_str:
+                image_meta, token_str = image_str.split(self.img_token)
+                H, W = tuple(map(int, image_meta.split(' ')))
+                token_ids = list(map(
+                    lambda x: int(x.split('>')[0]),
+                    token_str.split('<MM-Token-')[1:H*W+1],
+                ))
+                if len(force_text) > 0:
+                    image = self.detokenize_image([force_text.pop(0)], images, token_ids, (H, W))
+                else:
+                    image = self.detokenize_image(texts, images, token_ids, (H, W))
+                images.append(image)
+        return texts, images
+    
+    @torch.no_grad()
+    def tokenize_image(self, image):
+        assert hasattr(self, 'image_tokenizer'), 'Please call "init_vision" before that.'
+
+        image_str = self.som_token
+        image = self.image_transform(image)
+        assert image is not None, f'Unsupported image aspect ratio (max {self.transform_config.max_aspect_ratio}) or image resolution is too low (min {self.transform_config.min_short_size})'
+
+        image = image[None, ...].to(self.device, self.vision_dtype)
+        tokens = self.image_tokenizer.encode(image)
+        B, H, W = tokens.shape
+        tokens = tokens.view(B, -1).cpu().tolist()[0]
+        token_str = ''.join(map(lambda x: '<MM-Token-{token_id}>'.format(token_id=x), tokens))
+        image_str = f'{self.som_token}{H} {W}{self.img_token}{token_str}{self.eom_token}'
+        return image_str
+    
+    @torch.no_grad()
+    def detokenize_image(self, texts, images, token_ids, shape):
+        assert hasattr(self, 'image_tokenizer'), 'Please call "init_vision" before that.'
+        assert len(texts) == 1 and len(images) == 0, 'Only support one image per sample.'
+        H, W = shape
+        tokens = torch.tensor(token_ids, device=self.device, dtype=torch.long)
+        latents = self.image_tokenizer.decode(tokens, (1, H, W, self.encoder_config.codebook_dim))
+        upscale_factor = self.decoder_config.upscale_factor
+        latents = latents.reshape(*latents.shape[:2], -1).transpose(1, 2).contiguous()
+        image = self.decoder_pipe(
+            latents,
+            [texts[0]],
+            negative_prompt=[''],
+            height=H * upscale_factor, width=W * upscale_factor,
+            num_inference_steps=self.decoder_config.num_inference_steps, 
+            guidance_scale=1.0,
+            true_cfg_scale=self.decoder_config.cfg_scale,
+            true_cfg_scale_2=self.decoder_config.cfg_scale_2,
+        ).images[0]
+
+
+        return image
+    
+    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,
+        labels: Optional[torch.LongTensor] = 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,
+        num_logits_to_keep: int = 0,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        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
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        self.model.has_sliding_layers = False
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0]
+        hidden_states = hidden_states[:, -num_logits_to_keep:, :]
+        logits = hidden_states.new_full(
+            (*hidden_states.shape[:-1], self.config.vocab_size + self.config.mm_vocab_size), 
+            torch.finfo(hidden_states.dtype).min
+        )
+        if self.generation_mode == 'text':
+            logits[:, :, :self.config.vocab_size] = self.lm_head(hidden_states)
+        else:
+            logits[:, :, self.config.vocab_size:self.config.vocab_size + self.config.image_vocab_size] = self.mm_head(hidden_states)[:, :, :self.config.image_vocab_size]
+        
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # Upcast to float if we need to compute the loss to avoid potential precision issues
+            logits = logits.float()
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+AutoModel.register(XOmniConfig, XOmniModel)
+AutoModelForCausalLM.register(XOmniConfig, XOmniForCausalLM)