init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,10 @@
+.vscode
+__pycache__/
+*.pyc
+
+data
+checkpoint
+eval
+
+outputs
+wandb

app.py

@@ -0,0 +1,212 @@
+import gradio as gr
+import torch
+from src.transformer import SymmetricTransformer2DModel
+from src.pipeline import UnifiedPipeline
+from src.scheduler import Scheduler
+from torchvision import transforms
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer
+from diffusers import VQModel
+import os
+from PIL import Image
+import numpy as np
+
+
+def load_models(model_path="MeissonFlow/Meissonic", 
+            transformer_path="MeissonFlow/Muddit/1024",
+            device="cuda"):
+    model = SymmetricTransformer2DModel.from_pretrained(
+        transformer_path or model_path,
+        subfolder="transformer",
+    )
+    vq_model = VQModel.from_pretrained(model_path, subfolder="vqvae")
+    text_encoder = CLIPTextModelWithProjection.from_pretrained(model_path, subfolder="text_encoder")
+    tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+    scheduler = Scheduler.from_pretrained(model_path, subfolder="scheduler")
+
+    pipe = UnifiedPipeline(
+        vqvae=vq_model,
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
+        transformer=model,
+        scheduler=scheduler,
+    )
+    pipe.to(device)
+    return pipe
+
+# Load models (global variable to avoid reloading)
+pipe = load_models()
+
+# Common transform
+def get_transform(resolution):
+    return transforms.Compose([
+        transforms.Resize((resolution, resolution)),
+        transforms.ToTensor(),
+    ])
+
+# Image-to-Text Function
+def image_to_text(image, prompt, resolution=1024, steps=64, cfg=9.0):
+    try:
+        transform = get_transform(resolution)
+        
+        if image is not None:
+            pil_image = Image.fromarray(image.astype('uint8'), 'RGB') if isinstance(image, np.ndarray) else image
+            images = torch.stack([transform(pil_image)])
+            questions = [prompt] if prompt else ["Please describe this image."]
+        else:
+            images = None
+            questions = [prompt] if prompt else ["Please generate an image description."]
+
+        output = pipe(
+            prompt=questions,
+            image=images,
+            height=resolution,
+            width=resolution,
+            guidance_scale=cfg,
+            num_inference_steps=steps,
+            mask_token_embedding="MeissonFlow/Muddit",
+            generator=torch.manual_seed(42),
+        )
+
+        return output.prompts[0]
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Text-to-Image Function
+def text_to_image(prompt, negative_prompt, num_images=1, resolution=1024, steps=64, cfg=9.0):
+    try:
+        negative_prompt = negative_prompt or "worst quality, low quality, low res, blurry, distortion, watermark, logo, signature, text, jpeg artifacts, signature, sketch, duplicate, ugly, identifying mark"
+        
+        output = pipe(
+            prompt=[prompt]*num_images,
+            negative_prompt=[negative_prompt]*num_images,
+            height=resolution,
+            width=resolution,
+            guidance_scale=cfg,
+            num_inference_steps=steps,
+            mask_token_embedding="MeissonFlow/Muddit",
+            generator=torch.manual_seed(42),
+        )
+        
+        return output.images
+    
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return None
+
+# Create Gradio interface with Soft theme
+with gr.Blocks(theme=gr.themes.Soft(), title="Muddit Unifined Model") as demo:
+    gr.Markdown("# 🌌 Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Model.")
+    gr.Markdown(" Muddit is a unified discrete diffusion transformer that enables fast and parallel generation across both text and image modalities.")
+    
+    with gr.Tab("Image to Text"):
+        with gr.Row():
+            with gr.Column():
+                i2t_image_input = gr.Image(label="Upload Image", type="pil")
+                i2t_prompt_input = gr.Textbox(label="Prompt", value="Please describe this image.", placeholder="Enter your prompt here...")
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    i2t_resolution = gr.Slider(label="Resolution", minimum=256, maximum=1024, value=1024, step=64)
+                    i2t_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=64, step=1)
+                    i2t_cfg = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=9.0, step=0.5)
+                
+                i2t_submit_btn = gr.Button("Generate Description", variant="primary")
+            
+            with gr.Column():
+                i2t_output_text = gr.Textbox(label="Generated Description", interactive=False)
+                i2t_examples = gr.Examples(
+                    examples=[
+                        ["assets/man.jpg"],
+                        ["assets/tennis.jpg"],
+                        ["assets/pizza2.jpg"],
+                        ["assets/plane.jpg"],
+                        ["assets/zebra.jpg"],
+                        ["assets/building.jpg"],
+                        ["assets/flower.jpg"],
+                    ],
+                    inputs=[i2t_image_input],
+                    label="Example Inputs"
+                )
+    
+    with gr.Tab("VQA"):
+        with gr.Row():
+            with gr.Column():
+                vqa_image_input = gr.Image(label="Upload Image", type="pil")
+                vqa_prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your question here...")
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    vqa_resolution = gr.Slider(label="Resolution", minimum=256, maximum=1024, value=1024, step=64)
+                    vqa_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=64, step=1)
+                    vqa_cfg = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=9.0, step=0.5)
+                
+                vqa_submit_btn = gr.Button("Generate Answer", variant="primary")
+            
+            with gr.Column():
+                vqa_output_text = gr.Textbox(label="Generated Answer", interactive=False)
+                vqa_examples = gr.Examples(
+                    examples=[
+                        ["assets/kid.jpg", "What color is the kid's hair?"],
+                        ["assets/street.jpg", "Can someone legally walk across the street right now?"],
+                        ["assets/dog.jpg", "Where is the dog laying?"],
+                        ["assets/dog2.jpg", "What color is the toy the dog is holding?"],
+                        ["assets/pizza.jpg", "What food item is shown?"],
+                        ["assets/sheep.jpg", "How many sheep are pictured?"],
+                        ["assets/car.jpg", "Where are the cars?"],
+                    ],
+                    inputs=[vqa_image_input, vqa_prompt_input],
+                    label="Example Inputs"
+                )
+    
+    with gr.Tab("Text to Image"):
+        with gr.Row():
+            with gr.Column():
+                t2i_prompt_input = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate...")
+                t2i_negative_prompt = gr.Textbox(label="Negative Prompt", 
+                                               value="worst quality, low quality, low res, blurry, distortion, watermark, logo, signature, text, jpeg artifacts, signature, sketch, duplicate, ugly, identifying mark",
+                                               placeholder="What you don't want in the image...",
+                                               lines=5)
+                t2i_num_images = gr.Slider(label="Number of Images", minimum=1, maximum=4, value=1, step=1)
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    t2i_resolution = gr.Slider(label="Resolution", minimum=256, maximum=1024, value=1024, step=64)
+                    t2i_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=64, step=1)
+                    t2i_cfg = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=9.0, step=0.5)
+                
+                t2i_submit_btn = gr.Button("Generate Images", variant="primary")
+            
+            with gr.Column():
+                t2i_gallery = gr.Gallery(label="Generated Images")
+                t2i_examples = gr.Examples(
+                    examples=[
+                        ["A line art portrait showcasing a human figure with flowing, textured strokes"],
+                        ["A hyper realistic image of a chimpanzee with a glass-enclosed brain on his head, standing amidst lush, bioluminescent foliage in a vibrant futuristic forest"],
+                        ["A samurai in a stylized cyberpunk outfit adorned with intricate steampunk gear and floral accents, his Mandalorian armor gleaming under the backlighting"],
+                        ["A translucent, minimalist Porsche 911 GT3RS built from sleek carbon fiber, its aerodynamic body designed in the spirit of '60s Braun and modern Apple minimalism"],
+                        ["A realistic photograph of a ramadan tent shaped like a crescent moon under a velvety back sky studded with the milky way"],
+                        ["A portrait of John Lennon, captured in the gritty detail of line art"],
+                        ["In a world plunged into an unending darkness, remnants of fading starlight pierce through a heavy, smog-filled sky"]
+                    ],
+                    inputs=[t2i_prompt_input],
+                    label="Example Prompts"
+                )
+    
+    # Event handlers
+    i2t_submit_btn.click(
+        fn=image_to_text,
+        inputs=[i2t_image_input, i2t_prompt_input, i2t_resolution, i2t_steps, i2t_cfg],
+        outputs=i2t_output_text
+    )
+    
+    vqa_submit_btn.click(
+        fn=image_to_text,
+        inputs=[vqa_image_input, vqa_prompt_input, vqa_resolution, vqa_steps, vqa_cfg],
+        outputs=vqa_output_text
+    )
+    
+    t2i_submit_btn.click(
+        fn=text_to_image,
+        inputs=[t2i_prompt_input, t2i_negative_prompt, t2i_num_images, t2i_resolution, t2i_steps, t2i_cfg],
+        outputs=t2i_gallery
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,342 @@
+absl-py==2.1.0
+accelerate==0.34.2
+addict==2.4.0
+aiofiles==23.2.1
+aiohappyeyeballs==2.4.0
+aiohttp==3.10.5
+aiosignal==1.3.1
+albucore==0.0.19
+albumentations==1.4.20
+aliyun-python-sdk-core==2.16.0
+aliyun-python-sdk-kms==2.16.5
+all-clip==1.2.0
+aniso8601==9.0.1
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.9.0
+appdirs==1.4.4
+argon2-cffi==23.1.0
+argon2-cffi-bindings==21.2.0
+arrow==1.3.0
+asttokens==2.4.1
+async-lru==2.0.4
+async-timeout==4.0.3
+attrdict==2.0.1
+attrs==24.2.0
+av==14.0.1
+babel==2.16.0
+beartype==0.19.0
+beautifulsoup4==4.12.3
+bitsandbytes==0.43.3
+bleach==6.2.0
+blessed==1.20.0
+blinker==1.8.2
+boto3==1.35.92
+botocore==1.35.92
+braceexpand==0.1.7
+build==1.2.2.post1
+cachetools==5.5.2
+certifi==2024.8.30
+cffi==1.17.1
+charset-normalizer==2.0.12
+chumpy==0.70
+click==8.1.7
+clip-anytorch==2.6.0
+cloudpickle==3.1.0
+colorama==0.4.6
+comm==0.2.2
+contourpy==1.3.0
+cpm-kernels==1.0.11
+crcmod==1.7
+cryptography==44.0.3
+cycler==0.12.1
+Cython==3.0.12
+dashscope==1.22.2
+dataclasses==0.6
+datasets==3.6.0
+debugpy==1.8.8
+decorator==4.4.2
+decord==0.6.0
+deepspeed==0.16.2
+defusedxml==0.7.1
+diffusers==0.33.1
+dill==0.3.8
+distro==1.9.0
+docker-pycreds==0.4.0
+easydict==1.13
+einops==0.8.1
+eval_type_backport==0.2.0
+exceptiongroup==1.2.2
+executing==2.1.0
+ExifRead-nocycle==3.0.1
+fairscale==0.4.13
+fastapi==0.115.0
+fastdtw==0.3.4
+fastjsonschema==2.20.0
+ffmpy==0.4.0
+filelock==3.14.0
+filterpy==1.4.5
+fire==0.5.0
+flash_attn==2.7.4.post1
+Flask==3.0.3
+Flask-Cors==4.0.2
+Flask-RESTful==0.3.10
+flow-vis==0.1
+fonttools==4.53.1
+fqdn==1.5.1
+freetype-py==2.5.1
+frozenlist==1.4.1
+fsspec==2024.2.0
+ftfy==6.3.1
+func_timeout==4.3.5
+fvcore==0.1.5.post20221221
+gdown==5.2.0
+gitdb==4.0.11
+GitPython==3.1.43
+google-auth==2.38.0
+gpustat==1.1.1
+gradio_client==1.3.0
+grpcio==1.69.0
+h11==0.14.0
+h5py==3.12.1
+hjson==3.1.0
+httpcore==1.0.5
+httpx==0.28.1
+huggingface-hub==0.29.1
+hydra-core==1.3.2
+hydra-submitit-launcher==1.2.0
+idna==3.10
+imageio==2.35.1
+imageio-ffmpeg==0.5.1
+importlib_metadata==8.5.0
+importlib_resources==6.4.5
+imutils==0.5.4
+iopath==0.1.10
+ipykernel==6.29.5
+ipympl==0.9.4
+ipython==8.18.1
+ipython-genutils==0.2.0
+ipywidgets==8.1.5
+isoduration==20.11.0
+itsdangerous==2.2.0
+jaxtyping==0.2.36
+jedi==0.19.2
+Jinja2==3.1.3
+jiter==0.5.0
+jmespath==0.10.0
+joblib==1.4.2
+json-tricks==3.17.3
+json5==0.9.28
+jsonpointer==3.0.0
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
+jupyter==1.1.1
+jupyter-console==6.6.3
+jupyter-events==0.10.0
+jupyter-lsp==2.2.5
+jupyter_client==8.6.3
+jupyter_core==5.7.2
+jupyter_server==2.14.2
+jupyter_server_terminals==0.5.3
+jupyterlab==4.2.6
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.27.3
+jupyterlab_widgets==3.0.13
+kiwisolver==1.4.7
+kornia==0.7.3
+kornia_rs==0.1.8
+lazy_loader==0.4
+lightning-utilities==0.11.9
+llvmlite==0.43.0
+Markdown==3.7
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.7.0
+matplotlib-inline==0.1.7
+mdurl==0.1.2
+mediapy==1.2.2
+mistune==3.0.2
+mmcv==2.2.0
+mmengine==0.10.7
+mmpose==0.28.0
+model-index==0.1.11
+moviepy==1.0.3
+mpmath==1.3.0
+msgpack==1.1.0
+multidict==6.1.0
+multilingual-clip==1.0.10
+multiprocess==0.70.16
+munkres==1.1.4
+nbclient==0.10.0
+nbconvert==7.16.4
+nbformat==5.10.4
+nest-asyncio==1.6.0
+networkx==3.2.1
+ninja==1.11.1.3
+nltk==3.9.1
+notebook==7.2.2
+notebook_shim==0.2.4
+numba==0.60.0
+numpy==1.24.4
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-ml-py==12.560.30
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
+omegaconf==2.3.0
+open_clip_torch==2.29.0
+openai==1.47.0
+opencv-python==4.7.0.72
+opencv-python-headless==4.10.0.84
+opendatalab==0.0.10
+openmim==0.3.9
+ordered-set==4.1.0
+orjson==3.10.7
+oss2==2.17.0
+overrides==7.7.0
+packaging==24.1
+pandas==2.2.3
+pandocfilters==1.5.1
+parso==0.8.4
+peft==0.14.0
+pexpect==4.9.0
+Pillow==9.5.0
+pip-tools==7.4.1
+platformdirs==4.3.6
+plotly==5.24.1
+plyfile==1.1
+portalocker==2.10.1
+prodigyopt==1.0
+proglog==0.1.10
+prometheus_client==0.21.0
+prompt_toolkit==3.0.48
+protobuf==3.20.3
+psutil==6.0.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+py-cpuinfo==9.0.0
+pyarrow==20.0.0
+pyasn1==0.6.1
+pyasn1_modules==0.4.1
+pycocoevalcap==1.2
+pycocotools==2.0.8
+pycparser==2.22
+pycryptodome==3.22.0
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydub==0.25.1
+pyglet==1.5.27
+Pygments==2.18.0
+PyOpenGL==3.1.0
+pyparsing==3.1.4
+pyproject_hooks==1.2.0
+pyrender==0.1.45
+PySocks==1.7.1
+python-dateutil==2.9.0.post0
+python-json-logger==2.0.7
+python-multipart==0.0.10
+pytorch-lightning==2.5.0.post0
+pytz==2023.4
+PyWavelets==1.6.0
+PyYAML==6.0.2
+pyzmq==26.2.0
+qwen-vl-utils==0.0.10
+referencing==0.35.1
+regex==2024.9.11
+requests==2.32.3
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rich==13.4.2
+rpds-py==0.21.0
+rsa==4.9
+ruff==0.6.7
+s3transfer==0.10.4
+safetensors==0.4.5
+scikit-image==0.24.0
+scikit-learn==1.5.2
+scipy==1.10.1
+seaborn==0.13.2
+semantic-version==2.10.0
+Send2Trash==1.8.3
+sentence-transformers==2.7.0
+sentencepiece==0.1.99
+sentry-sdk==2.14.0
+setproctitle==1.3.3
+shellingham==1.5.4
+shortuuid==1.0.13
+simple-aesthetics-predictor==0.1.2
+six==1.16.0
+smmap==5.0.1
+smplx==0.1.28
+sniffio==1.3.1
+soupsieve==2.6
+stack-data==0.6.3
+starlette==0.38.6
+stringzilla==3.10.6
+submitit==1.5.2
+supervision==0.25.1
+SwissArmyTransformer==0.4.12
+sympy==1.13.1
+tabulate==0.9.0
+tenacity==9.0.0
+tensorboard==2.18.0
+tensorboard-data-server==0.7.2
+tensorboardX==2.6.2.2
+termcolor==2.5.0
+terminado==0.18.1
+threadpoolctl==3.5.0
+tifffile==2024.8.30
+timm==1.0.11
+tinycss2==1.4.0
+tokenizers==0.21.0
+tomesd==0.1.3
+tomli==2.1.0
+tomlkit==0.12.0
+torch==2.5.1
+torch-fidelity==0.3.0
+torchdata==0.11.0
+torchdiffeq==0.2.5
+torchgeometry==0.1.2
+torchmetrics==1.6.1
+torchsde==0.2.6
+torchvision==0.20.1
+tornado==6.4.1
+tqdm==4.67.1
+traitlets==5.14.3
+trampoline==0.1.2
+transformers==4.49.0
+trimesh==4.6.8
+triton==3.1.0
+tslearn==0.6.3
+typeguard==4.4.2
+typer==0.12.5
+types-python-dateutil==2.9.0.20241003
+typing_extensions==4.12.2
+tzdata==2024.1
+uri-template==1.3.0
+urllib3==1.26.20
+uvicorn==0.30.6
+wandb==0.17.5
+wcwidth==0.2.13
+webcolors==24.11.1
+webdataset==0.2.100
+webencodings==0.5.1
+websocket-client==1.8.0
+websockets==12.0
+Werkzeug==3.0.6
+widgetsnbextension==4.0.13
+wordcloud==1.9.4
+xtcocotools==1.14.3
+xxhash==3.5.0
+yacs==0.1.8
+yapf==0.43.0
+yarl==1.11.1
+zipp==3.20.2

src/pipeline.py

@@ -0,0 +1,684 @@
+# Copyright 2024 The HuggingFace Team and The MeissonFlow Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import sys
+from dataclasses import dataclass
+
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import PIL.Image
+import torch
+import PIL
+import numpy as np
+
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer, T5Tokenizer, T5EncoderModel
+from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
+from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models import VQModel
+from diffusers.utils import replace_example_docstring
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.utils import BaseOutput
+
+from src.scheduler import Scheduler
+from src.transformer import SymmetricTransformer2DModel
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+
+def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+    latent_image_ids = torch.zeros(height // 2, width // 2, 3)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]
+
+    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+    latent_image_ids = latent_image_ids.reshape(
+        latent_image_id_height * latent_image_id_width, latent_image_id_channels
+    )
+
+    return latent_image_ids.to(device=device, dtype=dtype)
+
+def dedup_consecutive_words(text: str) -> str:
+    """
+    >>> dedup_consecutive_words("hello hello world world world")
+    'hello world'
+    """
+    words = text.split()
+    if not words:
+        return text
+
+    out = [words[0]]
+    for w in words[1:]:
+        if w != out[-1]:
+            out.append(w)
+    return " ".join(out)
+
+def keep_upto_last_period(text: str) -> str:
+    """
+    Return the substring up to (and including) the last period-mark.
+    
+    The function searches first for the Chinese full stop “。”;
+    if none is found, it falls back to the ASCII dot “.”.
+    
+    Parameters
+    ----------
+    text : str
+        Input string.
+    
+    Returns
+    -------
+    str
+        Substring ending at the final period-mark.  If no period is present,
+        the original string is returned unchanged.
+    """
+    # Weired problem
+    text = text.replace("such is such", "").replace("is such is", "").replace("such is", "").replace("is such", "")
+    # Fallback to the ASCII period
+    idx = -1
+    if idx == -1:
+        idx = text.rfind(".")
+    # If still not found, return original text
+    if idx == -1:
+        return text
+    # Keep everything up to (and including) the last period
+    return text[:idx + 1]
+
+@dataclass
+class UnifiedPipelineOutput(BaseOutput):
+    """
+    Output class for image pipelines.
+
+    Args:
+        images (`List[PIL.Image.Image]` or `np.ndarray`)
+            List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+            num_channels)`.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+    prompts: List[str]
+
+
+class UnifiedPipeline(DiffusionPipeline):
+    image_processor: VaeImageProcessor
+    vqvae: VQModel
+    tokenizer: CLIPTokenizer
+    tokenizer_2: T5Tokenizer
+    text_encoder: CLIPTextModelWithProjection
+    text_encoder_2: T5EncoderModel
+    transformer: SymmetricTransformer2DModel
+    scheduler: Scheduler
+    model_cpu_offload_seq = "text_encoder->transformer->vqvae"
+
+    def __init__(
+        self,
+        vqvae: VQModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModelWithProjection,
+        transformer: SymmetricTransformer2DModel,
+        scheduler: Scheduler,
+        tokenizer_2: T5Tokenizer = None,
+        text_encoder_2: T5EncoderModel = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vqvae=vqvae,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vqvae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_normalize=False)
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Optional[Union[List[str], str]] = None,
+        height: Optional[int] = 1024,
+        width: Optional[int] = 1024,
+        image: Optional[Union[torch.Tensor, PIL.Image.Image]] = None,
+        num_inference_steps: int = 48,
+        guidance_scale: float = 9.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[torch.Generator] = None,
+        latents: Optional[torch.IntTensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_encoder_hidden_states: Optional[torch.Tensor] = None,
+        output_type = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
+        callback_steps: int = 1,
+        micro_conditioning_aesthetic_score: int = 6,
+        micro_conditioning_crop_coord: Tuple[int, int] = (0, 0),
+        temperature: Union[int, Tuple[int, int], List[int]] = (2, 0),
+        mask_token_embedding: Optional[str] = None,
+    ):
+        """
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            height (`int`, *optional*, defaults to `self.transformer.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 16):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 10.0):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.IntTensor`, *optional*):
+                Pre-generated tokens representing latent vectors in `self.vqvae`, to be used as inputs for image
+                gneration. If not provided, the starting latents will be completely masked.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument. A single vector from the
+                pooled and projected final hidden states.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                Pre-generated penultimate hidden states from the text encoder providing additional text conditioning.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+            negative_encoder_hidden_states (`torch.Tensor`, *optional*):
+                Analogous to `encoder_hidden_states` for the positive prompt.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            micro_conditioning_aesthetic_score (`int`, *optional*, defaults to 6):
+                The targeted aesthetic score according to the laion aesthetic classifier. See
+                https://laion.ai/blog/laion-aesthetics/ and the micro-conditioning section of
+                https://arxiv.org/abs/2307.01952.
+            micro_conditioning_crop_coord (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                The targeted height, width crop coordinates. See the micro-conditioning section of
+                https://arxiv.org/abs/2307.01952.
+            temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to (2, 0)):
+                Configures the temperature scheduler on `self.scheduler` see `Scheduler#set_timesteps`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.pipeline_utils.ImagePipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.pipeline_utils.ImagePipelineOutput`] is returned, otherwise a
+                `tuple` is returned where the first element is a list with the generated images.
+        """
+        if (prompt_embeds is not None and encoder_hidden_states is None) or (
+            prompt_embeds is None and encoder_hidden_states is not None
+        ):
+            raise ValueError("pass either both `prompt_embeds` and `encoder_hidden_states` or neither")
+
+        if (negative_prompt_embeds is not None and negative_encoder_hidden_states is None) or (
+            negative_prompt_embeds is None and negative_encoder_hidden_states is not None
+        ):
+            raise ValueError(
+                "pass either both `negatve_prompt_embeds` and `negative_encoder_hidden_states` or neither"
+            )
+        
+        if self.text_encoder_2 is not None:
+            self.text_encoder_2.to(self._execution_device)
+
+        text2image = image is None
+        image2text = image is not None
+
+        if image2text:
+            if self.text_encoder_2 is not None:
+                prompt = "<extra_id_0>" * 256
+                prompt = [prompt] * len(image)
+                
+                t5_mask_id = self.tokenizer_2.convert_tokens_to_ids("<extra_id_0>")
+                self.scheduler.config.mask_token_id = t5_mask_id
+            else:
+                mask_token = "<mask>"
+                self.tokenizer.add_tokens(mask_token, special_tokens=False)
+                clip_mask_id = self.tokenizer.convert_tokens_to_ids(mask_token)
+                self.text_encoder.resize_token_embeddings(len(self.tokenizer))
+                
+                if mask_token_embedding is not None:
+                    if mask_token_embedding.endswith(".pth"):
+                        mask_token_embedding = torch.load(mask_token_embedding)
+                    else:
+                        mask_token_embedding = os.path.dirname(mask_token_embedding)
+                        mask_token_embedding_path = os.path.join(mask_token_embedding, "mask_token_embedding.pth")
+                        assert os.path.exists(mask_token_embedding_path), f"{mask_token_embedding_path} doesn't exists!"
+                        mask_token_embedding = torch.load(mask_token_embedding_path)
+                        
+                    mask_token_embedding = mask_token_embedding.to(self._execution_device, dtype=self.text_encoder.dtype)
+                    self.text_encoder.get_input_embeddings().weight.data[clip_mask_id].copy_(mask_token_embedding)
+                
+                self.scheduler.config.mask_token_id = clip_mask_id
+                
+                input_ids = torch.ones(
+                    size=(len(image), self.tokenizer.model_max_length),
+                    dtype=torch.int64,
+                    device=self._execution_device
+                )
+                input_ids = input_ids * clip_mask_id
+                
+                question_len = []
+                if prompt is None:
+                    question_len = [0] * len(image)
+                elif isinstance(prompt, str):
+                    question_ids = torch.LongTensor([self.tokenizer.encode(prompt)])
+                    question_ids = question_ids.repeat(len(image), 1)
+                    
+                    q_len = len(question_ids[0]) - 1   # remove <eos> token
+                    question_len = [q_len] * len(image)
+                    
+                    input_ids[:, :q_len] = question_ids[:, :-1]
+                else: 
+                    assert isinstance(prompt, list), f"prompt must be None or str or list!"
+                    assert len(prompt) == len(image), f"VQA require equal num of images and prompts!"
+                    for i, p in enumerate(prompt):
+                        question_ids = torch.LongTensor([self.tokenizer.encode(p)])
+                        
+                        q_len = len(question_ids[0]) - 1
+                        question_len.append(q_len)
+                        
+                        input_ids[i, :q_len] = question_ids[0, :-1]
+        else:
+            self.scheduler.config.mask_token_id = self.transformer.config.vocab_size - 1
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        if image is not None:
+            batch_size = len(image)
+        else:
+            batch_size = len(prompt)
+
+        if height is None:
+            height = self.transformer.config.sample_size * self.vae_scale_factor
+
+        if width is None:
+            width = self.transformer.config.sample_size * self.vae_scale_factor
+
+        if isinstance(self.text_encoder, CLIPTextModelWithProjection):
+            text_encoder_type = "open_clip"
+        if isinstance(self.text_encoder_2, Gemma2Model):
+            text_encoder_type = "gemma"
+        
+        if prompt_embeds is None:
+            if text_encoder_type == "t5_clip":
+                if text2image:
+                    input_ids_clip = self.tokenizer(
+                        prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=77,
+                    ).input_ids.to(self._execution_device)
+                    outputs = self.text_encoder(input_ids_clip, return_dict=True, output_hidden_states=True)
+                    prompt_embeds = outputs.text_embeds
+                    
+                    input_ids_t5 = self.tokenizer_2(
+                        prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=256,
+                    ).input_ids.to(self._execution_device)
+                    
+                outputs_2 = self.text_encoder_2(input_ids_t5, return_dict=True, output_hidden_states=True)
+                encoder_hidden_states = outputs_2.last_hidden_state
+            elif text_encoder_type == "open_clip":
+                if text2image:
+                    input_ids = self.tokenizer(
+                        prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=77,
+                    ).input_ids.to(self._execution_device)
+                    
+                outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True)
+                prompt_embeds = outputs.text_embeds
+                encoder_hidden_states = outputs.hidden_states[-2]
+            elif text_encoder_type == "gemma":
+                if text2image:
+                    input_ids_clip = self.tokenizer(
+                        prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=77,
+                    ).input_ids.to(self._execution_device)
+                    outputs = self.text_encoder(input_ids_clip, return_dict=True, output_hidden_states=True)
+                    prompt_embeds = outputs.text_embeds
+                    
+                    input_ids_2 = self.tokenizer_2(
+                        prompt,
+                        truncation=True,
+                        padding="max_length",
+                        max_length=256,
+                        return_tensors="pt",
+                    ).input_ids.to(self._execution_device)
+                    
+                outputs_2 = self.text_encoder_2(input_ids_2, return_dict=True, output_hidden_states=True)
+                encoder_hidden_states = outputs_2.last_hidden_state
+
+        prompt_embeds = prompt_embeds.repeat(num_images_per_prompt, 1)
+        encoder_hidden_states = encoder_hidden_states.repeat(num_images_per_prompt, 1, 1)
+
+        if guidance_scale > 1.0 and text2image:
+            if negative_prompt_embeds is None:
+                if negative_prompt is None:
+                    negative_prompt = [""] * len(prompt)
+
+                if isinstance(negative_prompt, str):
+                    negative_prompt = [negative_prompt] * len(prompt)
+
+                if text_encoder_type == "t5_clip":                    
+                    input_ids = self.tokenizer(
+                        negative_prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=77,
+                    ).input_ids.to(self._execution_device)
+                    outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True)
+                    negative_prompt_embeds = outputs.text_embeds
+                    
+                    input_ids_2 = self.tokenizer_2(
+                        negative_prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=256,
+                    ).input_ids.to(self._execution_device)
+                    outputs_2 = self.text_encoder_2(input_ids_2, return_dict=True, output_hidden_states=True)
+                    negative_encoder_hidden_states = outputs_2.last_hidden_state
+                
+                elif text_encoder_type == "open_clip":
+                    input_ids = self.tokenizer(
+                        negative_prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=77,
+                    ).input_ids.to(self._execution_device)
+
+                    outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True)
+                
+                    negative_prompt_embeds = outputs.text_embeds
+                    negative_encoder_hidden_states = outputs.hidden_states[-2]
+                
+                elif text_encoder_type == "gemma":                    
+                    input_ids = self.tokenizer(
+                        negative_prompt,
+                        return_tensors="pt",
+                        padding="max_length",
+                        truncation=True,
+                        add_special_tokens=True,
+                        max_length=77,
+                    ).input_ids.to(self._execution_device)
+                    outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True)
+                    negative_prompt_embeds = outputs.text_embeds
+                    
+                    input_ids_2 = self.tokenizer_2(
+                        negative_prompt,
+                        truncation=True,
+                        padding="max_length",
+                        max_length=256,
+                        return_tensors="pt",
+                    ).input_ids.to(self._execution_device)
+                    outputs_2 = self.text_encoder_2(input_ids_2, return_dict=True, output_hidden_states=True)
+                    negative_encoder_hidden_states = outputs_2.last_hidden_state                 
+                          
+            negative_prompt_embeds = negative_prompt_embeds.repeat(num_images_per_prompt, 1)
+            negative_encoder_hidden_states = negative_encoder_hidden_states.repeat(num_images_per_prompt, 1, 1)
+
+            prompt_embeds = torch.concat([negative_prompt_embeds, prompt_embeds])
+            encoder_hidden_states = torch.concat([negative_encoder_hidden_states, encoder_hidden_states])
+
+        # Note that the micro conditionings _do_ flip the order of width, height for the original size
+        # and the crop coordinates. This is how it was done in the original code base
+        micro_conds = torch.tensor(
+            [
+                width,
+                height,
+                micro_conditioning_crop_coord[0],
+                micro_conditioning_crop_coord[1],
+                micro_conditioning_aesthetic_score,
+            ],
+            device=self._execution_device,
+            dtype=encoder_hidden_states.dtype,
+        )
+        micro_conds = micro_conds.unsqueeze(0)
+        micro_conds = micro_conds.expand(2 * batch_size if guidance_scale > 1.0 and text2image else batch_size, -1)
+
+        shape = (batch_size, height // self.vae_scale_factor, width // self.vae_scale_factor)
+
+        if latents is None and text2image:
+            latents = torch.full(
+                shape, self.scheduler.config.mask_token_id, dtype=torch.long, device=self._execution_device
+            )
+        elif image2text:
+            if text_encoder_type in ("t5_clip", "gemma"):
+                latents = input_ids_2 # [b, l]
+            else:
+                latents = input_ids
+
+        model_input = None
+
+        step_by_step = []
+
+        self.scheduler.set_timesteps(num_inference_steps, temperature, self._execution_device)
+        num_warmup_steps = len(self.scheduler.timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, timestep in enumerate(self.scheduler.timesteps):
+                if guidance_scale > 1.0 and text2image:
+                    model_input = torch.cat([latents] * 2)
+                    encoder_hidden_states = encoder_hidden_states
+                elif image2text:
+                    if model_input is None:
+                        model_input = self.vqvae.quantize(
+                            self.vqvae.encode(image.to(self._execution_device, dtype=self.vqvae.dtype)).latents
+                        )[2][2].reshape(batch_size, height // self.vae_scale_factor, width // self.vae_scale_factor)
+                    
+                    if text_encoder_type in ("t5_clip", "gemma"):
+                        outputs_t5 = self.text_encoder_2(latents, return_dict=True)
+                        encoder_hidden_states = outputs_t5.last_hidden_state
+                        
+                        batch_prompt = []
+                        for i in range(latents.size(0)):
+                            masked_prompt_input_id = latents[i].tolist()
+                            prompt = self.tokenizer_2.decode(masked_prompt_input_id, skip_special_tokens=True)
+                            batch_prompt.append(prompt)
+                        
+                        masked_prompt_input_ids_clip = self.tokenizer(
+                            batch_prompt,
+                            truncation=True,
+                            padding="max_length",
+                            max_length=77,
+                            return_tensors="pt"
+                        ).input_ids
+                        masked_prompt_input_ids_clip = masked_prompt_input_ids_clip.to(self._execution_device)
+                        outputs_clip = self.text_encoder(input_ids=masked_prompt_input_ids_clip, return_dict=True)
+                        prompt_embeds = outputs_clip.text_embeds
+                        
+                    else:
+                        outputs = self.text_encoder(latents, return_dict=True, output_hidden_states=True)
+                        prompt_embeds = outputs.text_embeds
+                        encoder_hidden_states = outputs.hidden_states[-2]
+                else:
+                    model_input = latents
+                    encoder_hidden_states = encoder_hidden_states
+                    
+                if height == 1024: #args.resolution == 1024:
+                    img_ids = _prepare_latent_image_ids(
+                        model_input.shape[0], 
+                        model_input.shape[-2],
+                        model_input.shape[-1],
+                        model_input.device,
+                        model_input.dtype
+                    )
+                else:
+                    img_ids = _prepare_latent_image_ids(
+                        model_input.shape[0],
+                        model_input.shape[-2],
+                        model_input.shape[-1],
+                        model_input.device,
+                        model_input.dtype
+                    )
+                txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(
+                    device=encoder_hidden_states.device,
+                    dtype=encoder_hidden_states.dtype
+                )
+                
+                # timestep_ = int(timestep / num_inference_steps * 1000)
+                model_output, encoder_hidden_states_tmp = self.transformer(
+                    hidden_states=model_input,
+                    micro_conds=micro_conds,
+                    pooled_projections=prompt_embeds,
+                    encoder_hidden_states=encoder_hidden_states,
+                    img_ids=img_ids,
+                    txt_ids=txt_ids,
+                    timestep=torch.tensor([timestep / num_inference_steps], device=model_input.device),
+                )
+
+                if image2text:
+                    encoder_hidden_states = encoder_hidden_states_tmp.clone()
+
+                if guidance_scale > 1.0 and text2image:
+                    uncond_logits, cond_logits = model_output.chunk(2)
+                    to_scheduler = uncond_logits + guidance_scale * (cond_logits - uncond_logits)
+                elif image2text:
+                    to_scheduler = encoder_hidden_states
+                else:
+                    to_scheduler = model_output
+
+                latents = self.scheduler.step(
+                    model_output=to_scheduler,
+                    timestep=timestep,
+                    sample=latents,
+                    generator=generator,
+                ).prev_sample
+
+                # this line will print the intermediate results of the image-to-text generation
+                # step_by_step.append(self.tokenizer.decode(latents[0].tolist(), skip_special_tokens=True))
+                
+                # this line will print the intermediate results of the text-to-image generation
+                # output = self.vqvae.decode(
+                #     latents,
+                #     force_not_quantize=True,
+                #     shape=(
+                #         batch_size,
+                #         height // self.vae_scale_factor,
+                #         width // self.vae_scale_factor,
+                #         self.vqvae.config.latent_channels,
+                #     ),
+                # ).sample.clip(0, 1)
+                # output = self.image_processor.postprocess(output, output_type)    # output is a list of PIL.Image, you need to save it.
+
+                if i == len(self.scheduler.timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, timestep, latents)
+
+        # with open("step_by_step.txt", "w") as file:
+        #     for prompt in step_by_step:
+        #         file.write(prompt + "\n")
+
+        if guidance_scale > 1.0 and text2image:
+            decoded_input_ids = encoder_hidden_states[encoder_hidden_states.shape[0] // 2:].argmax(-1)
+        else:
+            decoded_input_ids = encoder_hidden_states.argmax(-1)
+            
+        prompts = []
+        for i, prompt in enumerate(decoded_input_ids):
+            if image2text:
+                q_len = question_len[i]
+                prompt = self.tokenizer.decode(prompt.tolist()[q_len:], skip_special_tokens=True)        
+                prompts.append(keep_upto_last_period(dedup_consecutive_words(prompt)))
+            else:
+                prompts.append("Placeholder")
+            
+        if output_type == "latent":
+            output = latents
+        else:
+            needs_upcasting = self.vqvae.dtype == torch.float16 and self.vqvae.config.force_upcast
+
+            if needs_upcasting:
+                self.vqvae.float()
+
+            if text2image:
+                to_vqvae = latents
+            else:
+                to_vqvae = model_input
+                
+            output = self.vqvae.decode(
+                to_vqvae,
+                force_not_quantize=True,
+                shape=(
+                    batch_size,
+                    height // self.vae_scale_factor,
+                    width // self.vae_scale_factor,
+                    self.vqvae.config.latent_channels,
+                ),
+            ).sample.clip(0, 1)
+            output = self.image_processor.postprocess(output, output_type)
+
+            if needs_upcasting:
+                self.vqvae.half()
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (output,)
+
+        return UnifiedPipelineOutput(images=output, prompts=prompts)

src/scheduler.py

@@ -0,0 +1,175 @@
+# Copyright 2024 The HuggingFace Team and The MeissonFlow Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+def gumbel_noise(t, generator=None):
+    device = generator.device if generator is not None else t.device
+    noise = torch.zeros_like(t, device=device).uniform_(0, 1, generator=generator).to(t.device)
+    return -torch.log((-torch.log(noise.clamp(1e-20))).clamp(1e-20))
+
+
+def mask_by_random_topk(mask_len, probs, temperature=1.0, generator=None):
+    confidence = torch.log(probs.clamp(1e-20)) + temperature * gumbel_noise(probs, generator=generator)
+    sorted_confidence = torch.sort(confidence, dim=-1).values
+    cut_off = torch.gather(sorted_confidence, 1, mask_len.long())
+    masking = confidence < cut_off
+    return masking
+
+
+@dataclass
+class SchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
+            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
+            `pred_original_sample` can be used to preview progress or for guidance.
+    """
+
+    prev_sample: torch.Tensor
+    pred_original_sample: torch.Tensor = None
+
+
+class Scheduler(SchedulerMixin, ConfigMixin):
+    order = 1
+
+    temperatures: torch.Tensor
+
+    @register_to_config
+    def __init__(
+        self,
+        mask_token_id: int,
+        masking_schedule: str = "cosine",
+    ):
+        self.temperatures = None
+        self.timesteps = None
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        temperature: Union[int, Tuple[int, int], List[int]] = (2, 0),
+        device: Union[str, torch.device] = None,
+    ):
+        self.timesteps = torch.arange(num_inference_steps, device=device).flip(0)
+
+        if isinstance(temperature, (tuple, list)):
+            self.temperatures = torch.linspace(temperature[0], temperature[1], num_inference_steps, device=device)
+        else:
+            self.temperatures = torch.linspace(temperature, 0.01, num_inference_steps, device=device)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: torch.long,
+        sample: torch.LongTensor,
+        starting_mask_ratio: int = 1,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[SchedulerOutput, Tuple]:
+        two_dim_input = sample.ndim == 3 and model_output.ndim == 4
+
+        if two_dim_input:
+            batch_size, codebook_size, height, width = model_output.shape
+            sample = sample.reshape(batch_size, height * width)
+            model_output = model_output.reshape(batch_size, codebook_size, height * width).permute(0, 2, 1)
+
+        unknown_map = sample == self.config.mask_token_id
+
+        probs = model_output.softmax(dim=-1)
+
+        device = probs.device
+        probs_ = probs.to(generator.device) if generator is not None else probs  # handles when generator is on CPU
+        if probs_.device.type == "cpu" and probs_.dtype != torch.float32:
+            probs_ = probs_.float()  # multinomial is not implemented for cpu half precision
+        probs_ = probs_.reshape(-1, probs.size(-1))
+        pred_original_sample = torch.multinomial(probs_, 1, generator=generator).to(device=device)
+        pred_original_sample = pred_original_sample[:, 0].view(*probs.shape[:-1])
+        pred_original_sample = torch.where(unknown_map, pred_original_sample, sample)
+
+        if timestep == 0:
+            prev_sample = pred_original_sample
+        else:
+            seq_len = sample.shape[1]
+            step_idx = (self.timesteps == timestep).nonzero()
+            ratio = (step_idx + 1) / len(self.timesteps)
+
+            if self.config.masking_schedule == "cosine":
+                mask_ratio = torch.cos(ratio * math.pi / 2)
+            elif self.config.masking_schedule == "linear":
+                mask_ratio = 1 - ratio
+            else:
+                raise ValueError(f"unknown masking schedule {self.config.masking_schedule}")
+
+            mask_ratio = starting_mask_ratio * mask_ratio
+
+            mask_len = (seq_len * mask_ratio).floor()
+            # do not mask more than amount previously masked
+            mask_len = torch.min(unknown_map.sum(dim=-1, keepdim=True) - 1, mask_len)
+            # mask at least one
+            mask_len = torch.max(torch.tensor([1], device=model_output.device), mask_len)
+
+            selected_probs = torch.gather(probs, -1, pred_original_sample[:, :, None])[:, :, 0]
+            # Ignores the tokens given in the input by overwriting their confidence.
+            selected_probs = torch.where(unknown_map, selected_probs, torch.finfo(selected_probs.dtype).max)
+
+            masking = mask_by_random_topk(mask_len, selected_probs, self.temperatures[step_idx], generator)
+
+            # Masks tokens with lower confidence.
+            prev_sample = torch.where(masking, self.config.mask_token_id, pred_original_sample)
+
+        if two_dim_input:
+            prev_sample = prev_sample.reshape(batch_size, height, width)
+            pred_original_sample = pred_original_sample.reshape(batch_size, height, width)
+
+        if not return_dict:
+            return (prev_sample, pred_original_sample)
+
+        return SchedulerOutput(prev_sample, pred_original_sample)
+
+    def add_noise(self, sample, timesteps, generator=None):
+        step_idx = (self.timesteps == timesteps).nonzero()
+        ratio = (step_idx + 1) / len(self.timesteps)
+
+        if self.config.masking_schedule == "cosine":
+            mask_ratio = torch.cos(ratio * math.pi / 2)
+        elif self.config.masking_schedule == "linear":
+            mask_ratio = 1 - ratio
+        else:
+            raise ValueError(f"unknown masking schedule {self.config.masking_schedule}")
+
+        mask_indices = (
+            torch.rand(
+                sample.shape, device=generator.device if generator is not None else sample.device, generator=generator
+            ).to(sample.device)
+            < mask_ratio
+        )
+
+        masked_sample = sample.clone()
+
+        masked_sample[mask_indices] = self.config.mask_token_id
+
+        return masked_sample

src/transformer.py

@@ -0,0 +1,1459 @@
+# Copyright 2024 Black Forest Labs, The HuggingFace Team, The InstantX Team and The MeissonFlow Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Any, Dict, Optional, Tuple, Union, List
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention import FeedForward, BasicTransformerBlock, SkipFFTransformerBlock
+from diffusers.models.attention_processor import (
+    Attention,
+    AttentionProcessor,
+    FluxAttnProcessor2_0,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, GlobalResponseNorm, RMSNorm
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings,TimestepEmbedding, get_timestep_embedding #,FluxPosEmbed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput 
+from diffusers.models.resnet import Downsample2D, Upsample2D
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+
+def get_3d_rotary_pos_embed(
+    embed_dim, crops_coords, grid_size, temporal_size, theta: int = 10000, use_real: bool = True
+) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """
+    RoPE for video tokens with 3D structure.
+
+    Args:
+    embed_dim: (`int`):
+        The embedding dimension size, corresponding to hidden_size_head.
+    crops_coords (`Tuple[int]`):
+        The top-left and bottom-right coordinates of the crop.
+    grid_size (`Tuple[int]`):
+        The grid size of the spatial positional embedding (height, width).
+    temporal_size (`int`):
+        The size of the temporal dimension.
+    theta (`float`):
+        Scaling factor for frequency computation.
+    use_real (`bool`):
+        If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+
+    Returns:
+        `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.
+    """
+    start, stop = crops_coords
+    grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)
+    grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)
+    grid_t = np.linspace(0, temporal_size, temporal_size, endpoint=False, dtype=np.float32)
+
+    # Compute dimensions for each axis
+    dim_t = embed_dim // 4
+    dim_h = embed_dim // 8 * 3
+    dim_w = embed_dim // 8 * 3
+
+    # Temporal frequencies
+    freqs_t = 1.0 / (theta ** (torch.arange(0, dim_t, 2).float() / dim_t))
+    grid_t = torch.from_numpy(grid_t).float()
+    freqs_t = torch.einsum("n , f -> n f", grid_t, freqs_t)
+    freqs_t = freqs_t.repeat_interleave(2, dim=-1)
+
+    # Spatial frequencies for height and width
+    freqs_h = 1.0 / (theta ** (torch.arange(0, dim_h, 2).float() / dim_h))
+    freqs_w = 1.0 / (theta ** (torch.arange(0, dim_w, 2).float() / dim_w))
+    grid_h = torch.from_numpy(grid_h).float()
+    grid_w = torch.from_numpy(grid_w).float()
+    freqs_h = torch.einsum("n , f -> n f", grid_h, freqs_h)
+    freqs_w = torch.einsum("n , f -> n f", grid_w, freqs_w)
+    freqs_h = freqs_h.repeat_interleave(2, dim=-1)
+    freqs_w = freqs_w.repeat_interleave(2, dim=-1)
+
+    # Broadcast and concatenate tensors along specified dimension
+    def broadcast(tensors, dim=-1):
+        num_tensors = len(tensors)
+        shape_lens = {len(t.shape) for t in tensors}
+        assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+        shape_len = list(shape_lens)[0]
+        dim = (dim + shape_len) if dim < 0 else dim
+        dims = list(zip(*(list(t.shape) for t in tensors)))
+        expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+        assert all(
+            [*(len(set(t[1])) <= 2 for t in expandable_dims)]
+        ), "invalid dimensions for broadcastable concatenation"
+        max_dims = [(t[0], max(t[1])) for t in expandable_dims]
+        expanded_dims = [(t[0], (t[1],) * num_tensors) for t in max_dims]
+        expanded_dims.insert(dim, (dim, dims[dim]))
+        expandable_shapes = list(zip(*(t[1] for t in expanded_dims)))
+        tensors = [t[0].expand(*t[1]) for t in zip(tensors, expandable_shapes)]
+        return torch.cat(tensors, dim=dim)
+
+    freqs = broadcast((freqs_t[:, None, None, :], freqs_h[None, :, None, :], freqs_w[None, None, :, :]), dim=-1)
+
+    t, h, w, d = freqs.shape
+    freqs = freqs.view(t * h * w, d)
+
+    # Generate sine and cosine components
+    sin = freqs.sin()
+    cos = freqs.cos()
+
+    if use_real:
+        return cos, sin
+    else:
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs_cis
+
+
+def get_2d_rotary_pos_embed(embed_dim, crops_coords, grid_size, use_real=True):
+    """
+    RoPE for image tokens with 2d structure.
+
+    Args:
+    embed_dim: (`int`):
+        The embedding dimension size
+    crops_coords (`Tuple[int]`)
+        The top-left and bottom-right coordinates of the crop.
+    grid_size (`Tuple[int]`):
+        The grid size of the positional embedding.
+    use_real (`bool`):
+        If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+
+    Returns:
+        `torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`.
+    """
+    start, stop = crops_coords
+    grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)
+    grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)  # [2, W, H]
+
+    grid = grid.reshape([2, 1, *grid.shape[1:]])
+    pos_embed = get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=use_real)
+    return pos_embed
+
+
+def get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=False):
+    assert embed_dim % 4 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_rotary_pos_embed(
+        embed_dim // 2, grid[0].reshape(-1), use_real=use_real
+    )  # (H*W, D/2) if use_real else (H*W, D/4)
+    emb_w = get_1d_rotary_pos_embed(
+        embed_dim // 2, grid[1].reshape(-1), use_real=use_real
+    )  # (H*W, D/2) if use_real else (H*W, D/4)
+
+    if use_real:
+        cos = torch.cat([emb_h[0], emb_w[0]], dim=1)  # (H*W, D)
+        sin = torch.cat([emb_h[1], emb_w[1]], dim=1)  # (H*W, D)
+        return cos, sin
+    else:
+        emb = torch.cat([emb_h, emb_w], dim=1)  # (H*W, D/2)
+        return emb
+
+
+def get_2d_rotary_pos_embed_lumina(embed_dim, len_h, len_w, linear_factor=1.0, ntk_factor=1.0):
+    assert embed_dim % 4 == 0
+
+    emb_h = get_1d_rotary_pos_embed(
+        embed_dim // 2, len_h, linear_factor=linear_factor, ntk_factor=ntk_factor
+    )  # (H, D/4)
+    emb_w = get_1d_rotary_pos_embed(
+        embed_dim // 2, len_w, linear_factor=linear_factor, ntk_factor=ntk_factor
+    )  # (W, D/4)
+    emb_h = emb_h.view(len_h, 1, embed_dim // 4, 1).repeat(1, len_w, 1, 1)  # (H, W, D/4, 1)
+    emb_w = emb_w.view(1, len_w, embed_dim // 4, 1).repeat(len_h, 1, 1, 1)  # (H, W, D/4, 1)
+
+    emb = torch.cat([emb_h, emb_w], dim=-1).flatten(2)  # (H, W, D/2)
+    return emb
+
+
+def get_1d_rotary_pos_embed(
+    dim: int,
+    pos: Union[np.ndarray, int],
+    theta: float = 10000.0,
+    use_real=False,
+    linear_factor=1.0,
+    ntk_factor=1.0,
+    repeat_interleave_real=True,
+    freqs_dtype=torch.float32,  # torch.float32 (hunyuan, stable audio), torch.float64 (flux)
+):
+    """
+    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+    index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+    data type.
+
+    Args:
+        dim (`int`): Dimension of the frequency tensor.
+        pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+        theta (`float`, *optional*, defaults to 10000.0):
+            Scaling factor for frequency computation. Defaults to 10000.0.
+        use_real (`bool`, *optional*):
+            If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+        linear_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor for the context extrapolation. Defaults to 1.0.
+        ntk_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor for the NTK-Aware RoPE. Defaults to 1.0.
+        repeat_interleave_real (`bool`, *optional*, defaults to `True`):
+            If `True` and `use_real`, real part and imaginary part are each interleaved with themselves to reach `dim`.
+            Otherwise, they are concateanted with themselves.
+        freqs_dtype (`torch.float32` or `torch.float64`, *optional*, defaults to `torch.float32`):
+            the dtype of the frequency tensor.
+    Returns:
+        `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+    """
+    assert dim % 2 == 0
+
+    if isinstance(pos, int):
+        pos = np.arange(pos)
+    theta = theta * ntk_factor
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype)[: (dim // 2)] / dim)) / linear_factor  # [D/2]
+    t = torch.from_numpy(pos).to(freqs.device)  # type: ignore  # [S]
+    freqs = torch.outer(t, freqs)  # type: ignore   # [S, D/2]
+    if use_real and repeat_interleave_real:
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float()  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float()  # [S, D]
+        return freqs_cos, freqs_sin
+    elif use_real:
+        freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float()  # [S, D]
+        freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float()  # [S, D]
+        return freqs_cos, freqs_sin
+    else:
+        freqs_cis = torch.polar(torch.ones_like(freqs), freqs).float()  # complex64     # [S, D/2]
+        return freqs_cis
+
+
+class FluxPosEmbed(nn.Module):
+    # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
+    def __init__(self, theta: int, axes_dim: List[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        n_axes = ids.shape[-1]
+        cos_out = []
+        sin_out = []
+        pos = ids.squeeze().float().cpu().numpy()
+        is_mps = ids.device.type == "mps"
+        freqs_dtype = torch.float32 if is_mps else torch.float64
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i], pos[:, i], repeat_interleave_real=True, use_real=True, freqs_dtype=freqs_dtype
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+
+class FusedFluxAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "FusedFluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        qkv = attn.to_qkv(hidden_states)
+        split_size = qkv.shape[-1] // 3
+        query, key, value = torch.split(qkv, split_size, dim=-1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        # `context` projections.
+        if encoder_hidden_states is not None:
+            encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
+            split_size = encoder_qkv.shape[-1] // 3
+            (
+                encoder_hidden_states_query_proj,
+                encoder_hidden_states_key_proj,
+                encoder_hidden_states_value_proj,
+            ) = torch.split(encoder_qkv, split_size, dim=-1)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states
+
+
+
+@maybe_allow_in_graph
+class  SingleTransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, mlp_ratio=4.0):
+        super().__init__()
+        self.mlp_hidden_dim = int(dim * mlp_ratio)
+
+        self.norm = AdaLayerNormZeroSingle(dim)
+        self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
+        self.act_mlp = nn.GELU(approximate="tanh")
+        self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
+
+        processor = FluxAttnProcessor2_0()
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            bias=True,
+            processor=processor,
+            qk_norm="rms_norm",
+            eps=1e-6,
+            pre_only=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+    ):
+        residual = hidden_states
+        norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+
+        attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        gate = gate.unsqueeze(1)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return hidden_states
+
+@maybe_allow_in_graph
+class TransformerBlock(nn.Module):
+    r"""
+    A Transformer block following the MMDiT architecture, introduced in Stable Diffusion 3.
+
+    Reference: https://arxiv.org/abs/2403.03206
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        context_pre_only (`bool`): Boolean to determine if we should add some blocks associated with the
+            processing of `context` conditions.
+    """
+
+    def __init__(self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6):
+        super().__init__()
+
+        self.norm1 = AdaLayerNormZero(dim)
+
+        self.norm1_context = AdaLayerNormZero(dim)
+
+        if hasattr(F, "scaled_dot_product_attention"):
+            processor = FluxAttnProcessor2_0()
+        else:
+            raise ValueError(
+                "The current PyTorch version does not support the `scaled_dot_product_attention` function."
+            )
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            added_kv_proj_dim=dim,
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            processor=processor,
+            qk_norm=qk_norm,
+            eps=eps,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        # self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="swiglu")
+
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        # self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="swiglu")
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor,
+        temb: torch.FloatTensor,
+        image_rotary_emb=None,
+    ):
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)
+
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
+            encoder_hidden_states, emb=temb
+        )
+        # Attention.
+        attn_output, context_attn_output = self.attn(
+            hidden_states=norm_hidden_states,
+            encoder_hidden_states=norm_encoder_hidden_states,
+            image_rotary_emb=image_rotary_emb,
+        )
+
+        # Process attention outputs for the `hidden_states`.
+        attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden_states = self.norm2(hidden_states)
+        norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = hidden_states + ff_output
+
+        # Process attention outputs for the `encoder_hidden_states`.
+
+        context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output
+        encoder_hidden_states = encoder_hidden_states + context_attn_output
+
+        norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+        norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class UVit2DConvEmbed(nn.Module):
+    def __init__(self, in_channels, block_out_channels, vocab_size, elementwise_affine, eps, bias):
+        super().__init__()
+        self.embeddings = nn.Embedding(vocab_size, in_channels)
+        self.layer_norm = RMSNorm(in_channels, eps, elementwise_affine)
+        self.conv = nn.Conv2d(in_channels, block_out_channels, kernel_size=1, bias=bias)
+
+    def forward(self, input_ids):
+        embeddings = self.embeddings(input_ids)
+        embeddings = self.layer_norm(embeddings)
+        embeddings = embeddings.permute(0, 3, 1, 2)
+        embeddings = self.conv(embeddings)
+        return embeddings
+
+class ConvMlmLayer(nn.Module):
+    def __init__(
+        self,
+        block_out_channels: int,
+        in_channels: int,
+        use_bias: bool,
+        ln_elementwise_affine: bool,
+        layer_norm_eps: float,
+        codebook_size: int,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2d(block_out_channels, in_channels, kernel_size=1, bias=use_bias)
+        self.layer_norm = RMSNorm(in_channels, layer_norm_eps, ln_elementwise_affine)
+        self.conv2 = nn.Conv2d(in_channels, codebook_size, kernel_size=1, bias=use_bias)
+
+    def forward(self, hidden_states):
+        hidden_states = self.conv1(hidden_states)
+        hidden_states = self.layer_norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+        logits = self.conv2(hidden_states)
+        return logits
+
+class SwiGLU(nn.Module):
+    r"""
+    A [variant](https://arxiv.org/abs/2002.05202) of the gated linear unit activation function. It's similar to `GEGLU`
+    but uses SiLU / Swish instead of GeLU.
+
+    Parameters:
+        dim_in (`int`): The number of channels in the input.
+        dim_out (`int`): The number of channels in the output.
+        bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+    """
+
+    def __init__(self, dim_in: int, dim_out: int, bias: bool = True):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)
+        self.activation = nn.SiLU()
+
+    def forward(self, hidden_states):
+        hidden_states = self.proj(hidden_states)
+        hidden_states, gate = hidden_states.chunk(2, dim=-1)
+        return hidden_states * self.activation(gate)
+
+class ConvNextBlock(nn.Module):
+    def __init__(
+        self, channels, layer_norm_eps, ln_elementwise_affine, use_bias, hidden_dropout, hidden_size, res_ffn_factor=4
+    ):
+        super().__init__()
+        self.depthwise = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=3,
+            padding=1,
+            groups=channels,
+            bias=use_bias,
+        )
+        self.norm = RMSNorm(channels, layer_norm_eps, ln_elementwise_affine)
+        self.channelwise_linear_1 = nn.Linear(channels, int(channels * res_ffn_factor), bias=use_bias)
+        self.channelwise_act = nn.GELU()
+        self.channelwise_norm = GlobalResponseNorm(int(channels * res_ffn_factor))
+        self.channelwise_linear_2 = nn.Linear(int(channels * res_ffn_factor), channels, bias=use_bias)
+        self.channelwise_dropout = nn.Dropout(hidden_dropout)
+        self.cond_embeds_mapper = nn.Linear(hidden_size, channels * 2, use_bias)
+
+    def forward(self, x, cond_embeds):
+        x_res = x
+
+        x = self.depthwise(x)
+
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+
+        x = self.channelwise_linear_1(x)
+        x = self.channelwise_act(x)
+        x = self.channelwise_norm(x)
+        x = self.channelwise_linear_2(x)
+        x = self.channelwise_dropout(x)
+
+        x = x.permute(0, 3, 1, 2)
+
+        x = x + x_res
+
+        scale, shift = self.cond_embeds_mapper(F.silu(cond_embeds)).chunk(2, dim=1)
+        x = x * (1 + scale[:, :, None, None]) + shift[:, :, None, None]
+
+        return x
+
+class Simple_UVitBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        ln_elementwise_affine,
+        layer_norm_eps,
+        use_bias,
+        downsample: bool,
+        upsample: bool,
+    ):
+        super().__init__()
+
+        if downsample:
+            self.downsample = Downsample2D(
+                channels,
+                use_conv=True,
+                padding=0,
+                name="Conv2d_0",
+                kernel_size=2,
+                norm_type="rms_norm",
+                eps=layer_norm_eps,
+                elementwise_affine=ln_elementwise_affine,
+                bias=use_bias,
+            )
+        else:
+            self.downsample = None
+
+        if upsample:
+            self.upsample = Upsample2D(
+                channels,
+                use_conv_transpose=True,
+                kernel_size=2,
+                padding=0,
+                name="conv",
+                norm_type="rms_norm",
+                eps=layer_norm_eps,
+                elementwise_affine=ln_elementwise_affine,
+                bias=use_bias,
+                interpolate=False,
+            )
+        else:
+            self.upsample = None
+
+    def forward(self, x):
+        # print("before,", x.shape)
+        if self.downsample is not None:
+            # print('downsample')
+            x = self.downsample(x)
+
+        if self.upsample is not None:
+            # print('upsample')
+            x = self.upsample(x)
+        # print("after,", x.shape)
+        return x
+
+class Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    """
+    The Transformer model introduced in Flux.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Parameters:
+        patch_size (`int`): Patch size to turn the input data into small patches.
+        in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
+        num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
+        num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
+        attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
+        num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
+        joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
+        guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
+    """
+
+    _supports_gradient_checkpointing = False #True 
+    # Due to NotImplementedError: DDPOptimizer backend: Found a higher order op in the graph. This is not supported. Please turn off DDP optimizer using torch._dynamo.config.optimize_ddp=False. Note that this can cause performance degradation because there will be one bucket for the entire Dynamo graph. 
+    # Please refer to this issue - https://github.com/pytorch/pytorch/issues/104674.
+    _no_split_modules = ["TransformerBlock", "SingleTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        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, # unused in our implementation
+        axes_dims_rope: Tuple[int] = (16, 56, 56),
+        vocab_size: int = 8256,
+        codebook_size: int = 8192,
+        downsample: bool = False,
+        upsample: bool = False,
+    ):
+        super().__init__()
+        self.out_channels = in_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim 
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
+     
+        self.transformer_blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_single_layers)
+            ]
+        )
+
+
+        self.gradient_checkpointing = False
+
+        in_channels_embed = self.inner_dim 
+        ln_elementwise_affine = True
+        layer_norm_eps = 1e-06
+        use_bias = False
+        micro_cond_embed_dim = 1280
+        self.embed = UVit2DConvEmbed(
+            in_channels_embed, self.inner_dim, self.config.vocab_size, ln_elementwise_affine, layer_norm_eps, use_bias
+        )
+        self.mlm_layer = ConvMlmLayer(
+            self.inner_dim, in_channels_embed, use_bias, ln_elementwise_affine, layer_norm_eps, self.config.codebook_size
+        )
+        self.cond_embed = TimestepEmbedding( 
+            micro_cond_embed_dim + self.config.pooled_projection_dim, self.inner_dim, sample_proj_bias=use_bias
+        )
+        self.encoder_proj_layer_norm = RMSNorm(self.inner_dim, layer_norm_eps, ln_elementwise_affine)
+        self.project_to_hidden_norm = RMSNorm(in_channels_embed, layer_norm_eps, ln_elementwise_affine)
+        self.project_to_hidden = nn.Linear(in_channels_embed, self.inner_dim, bias=use_bias)
+        self.project_from_hidden_norm = RMSNorm(self.inner_dim, layer_norm_eps, ln_elementwise_affine)
+        self.project_from_hidden = nn.Linear(self.inner_dim, in_channels_embed, bias=use_bias)
+        
+        self.down_block = Simple_UVitBlock(
+            self.inner_dim, 
+            ln_elementwise_affine,
+            layer_norm_eps,
+            use_bias,
+            downsample,
+            False,
+        )
+        self.up_block = Simple_UVitBlock(
+            self.inner_dim, #block_out_channels,
+            ln_elementwise_affine,
+            layer_norm_eps,
+            use_bias,
+            False,
+            upsample=upsample,
+        )
+       
+        # self.fuse_qkv_projections()
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedFluxAttnProcessor2_0
+    def fuse_qkv_projections(self):
+        """
+        Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
+        are fused. For cross-attention modules, key and value projection matrices are fused.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+        """
+        self.original_attn_processors = None
+
+        for _, attn_processor in self.attn_processors.items():
+            if "Added" in str(attn_processor.__class__.__name__):
+                raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
+
+        self.original_attn_processors = self.attn_processors
+
+        for module in self.modules():
+            if isinstance(module, Attention):
+                module.fuse_projections(fuse=True)
+
+        self.set_attn_processor(FusedFluxAttnProcessor2_0())
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
+    def unfuse_qkv_projections(self):
+        """Disables the fused QKV projection if enabled.
+
+        <Tip warning={true}>
+
+        This API is 🧪 experimental.
+
+        </Tip>
+
+        """
+        if self.original_attn_processors is not None:
+            self.set_attn_processor(self.original_attn_processors)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    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,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples= None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        micro_conds: 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.
+        """
+        micro_cond_encode_dim = 256 # same as self.config.micro_cond_encode_dim = 256 from amused
+        micro_cond_embeds = get_timestep_embedding(
+            micro_conds.flatten(), micro_cond_encode_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        ) 
+        micro_cond_embeds = micro_cond_embeds.reshape((hidden_states.shape[0], -1)) 
+
+        pooled_projections = torch.cat([pooled_projections, micro_cond_embeds], dim=1)
+        pooled_projections = pooled_projections.to(dtype=self.dtype)
+        pooled_projections = self.cond_embed(pooled_projections).to(encoder_hidden_states.dtype)    
+       
+
+        hidden_states = self.embed(hidden_states) 
+
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states) 
+        encoder_hidden_states = self.encoder_proj_layer_norm(encoder_hidden_states)
+        hidden_states = self.down_block(hidden_states)
+
+        batch_size, channels, height, width = hidden_states.shape
+        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+        hidden_states = self.project_to_hidden_norm(hidden_states) 
+        hidden_states = self.project_to_hidden(hidden_states)
+
+       
+        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."
+                )
+
+        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)
+        ) 
+
+        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) 
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,  
+                    image_rotary_emb=image_rotary_emb,
+                )
+                
+
+            # 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))
+                hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+            # 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.project_from_hidden_norm(hidden_states)
+        hidden_states = self.project_from_hidden(hidden_states)
+       
+
+        hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+
+        hidden_states = self.up_block(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+        
+        output = self.mlm_layer(hidden_states)
+        # self.unfuse_qkv_projections()
+        if not return_dict:
+            return (output,)
+
+    
+        return output
+    
+
+class SymmetricTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    """
+    The Transformer model introduced in Flux.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Parameters:
+        patch_size (`int`): Patch size to turn the input data into small patches.
+        in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
+        num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
+        num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
+        attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
+        num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
+        joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
+        guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
+    """
+
+    _supports_gradient_checkpointing = False #True 
+    # Due to NotImplementedError: DDPOptimizer backend: Found a higher order op in the graph. This is not supported. Please turn off DDP optimizer using torch._dynamo.config.optimize_ddp=False. Note that this can cause performance degradation because there will be one bucket for the entire Dynamo graph. 
+    # Please refer to this issue - https://github.com/pytorch/pytorch/issues/104674.
+    _no_split_modules = ["TransformerBlock", "SingleTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        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, # unused in our implementation
+        axes_dims_rope: Tuple[int] = (16, 56, 56),
+        vocab_size: int = 8256,
+        codebook_size: int = 8192,
+        tokenizer_vocab_size: Optional[int] = None,
+        t5_dim: Optional[int] = None,
+        downsample: bool = False,
+        upsample: bool = False,
+    ):
+        super().__init__()
+        self.out_channels = in_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim 
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=self.config.pooled_projection_dim
+        )
+
+        if t5_dim is not None:
+            self.adapter = nn.Sequential(
+                nn.LayerNorm(t5_dim, elementwise_affine=False, eps=1e-6),
+                nn.Linear(t5_dim, self.config.joint_attention_dim, bias=False)
+            )
+        else:
+            self.adapter = None
+
+        self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
+     
+        self.transformer_blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_single_layers)
+            ]
+        )
+
+        self.gradient_checkpointing = False
+
+        in_channels_embed = self.inner_dim 
+        ln_elementwise_affine = True
+        layer_norm_eps = 1e-06
+        use_bias = False
+        micro_cond_embed_dim = 1280
+        self.embed = UVit2DConvEmbed(
+            in_channels_embed, self.inner_dim, self.config.vocab_size, ln_elementwise_affine, layer_norm_eps, use_bias
+        )
+        self.mlm_layer = ConvMlmLayer(
+            self.inner_dim, in_channels_embed, use_bias, ln_elementwise_affine, layer_norm_eps, self.config.codebook_size
+        )
+        self.cond_embed = TimestepEmbedding( 
+            micro_cond_embed_dim + self.config.pooled_projection_dim, self.inner_dim, sample_proj_bias=use_bias
+        )
+        self.encoder_proj_layer_norm = RMSNorm(self.inner_dim, layer_norm_eps, ln_elementwise_affine)
+        self.project_to_hidden_norm = RMSNorm(in_channels_embed, layer_norm_eps, ln_elementwise_affine)
+        self.project_to_hidden = nn.Linear(in_channels_embed, self.inner_dim, bias=use_bias)
+        self.project_from_hidden_norm = RMSNorm(self.inner_dim, layer_norm_eps, ln_elementwise_affine)
+        self.project_from_hidden = nn.Linear(self.inner_dim, in_channels_embed, bias=use_bias)
+        
+        self.down_block = Simple_UVitBlock(
+            self.inner_dim, 
+            ln_elementwise_affine,
+            layer_norm_eps,
+            use_bias,
+            downsample,
+            False,
+        )
+        self.up_block = Simple_UVitBlock(
+            self.inner_dim,
+            ln_elementwise_affine,
+            layer_norm_eps,
+            use_bias,
+            False,
+            upsample=upsample,
+        )
+        
+        if tokenizer_vocab_size is not None:
+            self.text_decoder = nn.Sequential(
+                nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6),
+                nn.Linear(self.inner_dim, tokenizer_vocab_size, bias=use_bias)
+            )
+        else:
+            self.text_decoder = None
+    
+
+    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,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples= None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        micro_conds: 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.
+        """
+        micro_cond_encode_dim = 256 # same as self.config.micro_cond_encode_dim = 256 from amused
+        micro_cond_embeds = get_timestep_embedding(
+            micro_conds.flatten(), micro_cond_encode_dim, flip_sin_to_cos=True, downscale_freq_shift=0
+        ) 
+        micro_cond_embeds = micro_cond_embeds.reshape((hidden_states.shape[0], -1)) 
+
+        if self.adapter is not None:
+            encoder_hidden_states = self.adapter(encoder_hidden_states)
+
+        pooled_projections = torch.cat([pooled_projections, micro_cond_embeds], dim=1)
+        pooled_projections = pooled_projections.to(dtype=self.dtype)
+        pooled_projections = self.cond_embed(pooled_projections).to(encoder_hidden_states.dtype)    
+       
+        hidden_states = self.embed(hidden_states) 
+
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states) 
+        encoder_hidden_states = self.encoder_proj_layer_norm(encoder_hidden_states)
+        hidden_states = self.down_block(hidden_states)
+
+        batch_size, channels, height, width = hidden_states.shape
+        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, height * width, channels)
+        hidden_states = self.project_to_hidden_norm(hidden_states) 
+        hidden_states = self.project_to_hidden(hidden_states)
+
+       
+        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."
+                )
+
+        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)
+        ) 
+
+        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) 
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,  
+                    image_rotary_emb=image_rotary_emb,
+                )
+                
+
+            # 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))
+                hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+            # 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]
+                )
+
+        encoder_hidden_states = hidden_states[:, :encoder_hidden_states.shape[1], ...]
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1]:, ...]
+       
+        if self.text_decoder is not None:
+            encoder_hidden_states = self.text_decoder(encoder_hidden_states)
+            
+        hidden_states = self.project_from_hidden_norm(hidden_states)
+        hidden_states = self.project_from_hidden(hidden_states)
+
+        hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+
+        hidden_states = self.up_block(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+        
+        output = self.mlm_layer(hidden_states)
+        # self.unfuse_qkv_projections()
+        if not return_dict:
+            return (output, encoder_hidden_states)
+
+    
+        return output, encoder_hidden_states    # [b, l, tokenizer_vocab_size]