Upload folder using huggingface_hub

README.md

@@ -1,13 +1,100 @@
 ---
-title: Muddit
-emoji: 🌖
-colorFrom: red
+title: Muddit Interface
+emoji: 🎨
+colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 5.38.2
+sdk_version: 4.0.0
 app_file: app.py
 pinned: false
 license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🎨 Muddit Interface
+
+A unified model interface for **Text-to-Image generation** and **Visual Question Answering (VQA)** powered by advanced transformer architectures.
+
+## ✨ Features
+
+### 🖼️ Text-to-Image Generation
+- Generate high-quality images from detailed text descriptions
+- Customizable parameters (resolution, inference steps, CFG scale, seed)
+- Support for negative prompts to avoid unwanted elements
+- Real-time generation with progress tracking
+
+### ❓ Visual Question Answering
+- Upload images and ask natural language questions
+- Get detailed descriptions and answers about image content
+- Support for various question types (counting, description, identification)
+- Advanced visual understanding capabilities
+
+## 🚀 How to Use
+
+### Text-to-Image
+1. Go to the **"🖼️ Text-to-Image"** tab
+2. Enter your text description in the **Prompt** field
+3. Optionally add a **Negative Prompt** to exclude unwanted elements
+4. Adjust parameters as needed:
+   - **Width/Height**: Image resolution (256-1024px)
+   - **Inference Steps**: Quality vs speed (1-100)
+   - **CFG Scale**: Prompt adherence (1.0-20.0)
+   - **Seed**: For reproducible results
+5. Click **"🎨 Generate Image"**
+
+### Visual Question Answering
+1. Go to the **"❓ Visual Question Answering"** tab
+2. **Upload an image** using the image input
+3. **Ask a question** about the image
+4. Adjust processing parameters if needed
+5. Click **"🤔 Ask Question"** to get an answer
+
+## 📝 Example Prompts
+
+### Text-to-Image Examples:
+- "A majestic night sky awash with billowing clouds, sparkling with a million twinkling stars"
+- "A hyper realistic image of a chimpanzee with a glass-enclosed brain on his head, standing amidst lush, bioluminescent foliage"
+- "A samurai in a stylized cyberpunk outfit adorned with intricate steampunk gear and floral accents"
+
+### VQA Examples:
+- "What objects do you see in this image?"
+- "How many people are in the picture?"
+- "What is the main subject of this image?"
+- "Describe the scene in detail"
+- "What colors dominate this image?"
+
+## 🛠️ Technical Details
+
+- **Architecture**: Unified transformer-based model
+- **Text Encoder**: CLIP for text understanding
+- **Vision Encoder**: VQ-VAE for image processing
+- **Generation**: Advanced diffusion-based synthesis
+- **VQA**: Multimodal understanding with attention mechanisms
+
+## ⚙️ Parameters Guide
+
+| Parameter | Description | Recommended Range |
+|-----------|-------------|-------------------|
+| **Inference Steps** | More steps = higher quality, slower generation | 20-64 |
+| **CFG Scale** | How closely to follow the prompt | 7.0-12.0 |
+| **Resolution** | Output image size | 512x512 to 1024x1024 |
+| **Seed** | For reproducible results | Any integer or -1 for random |
+
+## 🎯 Use Cases
+
+- **Creative Content**: Generate artwork, illustrations, concepts
+- **Visual Analysis**: Analyze and understand image content
+- **Education**: Learn about visual AI and multimodal models
+- **Research**: Explore capabilities of unified vision-language models
+- **Accessibility**: Describe images for visually impaired users
+
+## 📄 License
+
+This project is licensed under the Apache 2.0 License.
+
+## 🤝 Contributing
+
+Feedback and contributions are welcome! Please feel free to submit issues or pull requests.
+
+---
+
+*Powered by Gradio and Hugging Face Spaces* 🤗 

app.py

@@ -0,0 +1,322 @@
+import sys
+import os
+import warnings
+import tempfile
+
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+)
+from diffusers.models.autoencoders.vq_model import VQModel
+
+from src.transformer import SymmetricTransformer2DModel
+from src.pipeline import UnifiedPipeline
+from src.scheduler import Scheduler
+from train.trainer_utils import load_images_to_tensor
+
+# Suppress FutureWarnings to reduce clutter
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+# Set Gradio temp directory to a writable location
+def setup_gradio_temp_dir():
+    """Setup a writable temp directory for Gradio with fallback options"""
+    possible_dirs = [
+        os.path.join(os.getcwd(), "gradio_tmp"),  # Project directory
+        os.path.join(os.path.expanduser("~"), ".gradio_tmp"),  # Home directory
+        tempfile.mkdtemp(prefix="gradio_")  # System temp with unique name
+    ]
+    
+    for temp_dir in possible_dirs:
+        try:
+            os.makedirs(temp_dir, exist_ok=True)
+            # Test write permission
+            test_file = os.path.join(temp_dir, "test_write.tmp")
+            with open(test_file, "w") as f:
+                f.write("test")
+            os.remove(test_file)
+            
+            os.environ["GRADIO_TEMP_DIR"] = temp_dir
+            print(f"✅ Gradio temp directory set to: {temp_dir}")
+            return temp_dir
+        except (PermissionError, OSError) as e:
+            print(f"⚠️  Cannot use {temp_dir}: {e}")
+            continue
+    
+    raise RuntimeError("Could not find a writable directory for Gradio temp files")
+
+setup_gradio_temp_dir()
+
+
+
+class MudditInterface:
+    def __init__(self, model_path="MeissonFlow/Meissonic", transformer_path="QingyuShi/Muddit"):
+        if torch.cuda.is_available():
+            device = "cuda"
+        else:
+            device = "cpu"
+        self.device = device
+        self.model_path = model_path
+        self.transformer_path = transformer_path or model_path
+        
+        print("Loading models...")
+        self.load_models()
+        print("Models loaded successfully!")
+    
+    def load_models(self):
+        """Load all required models"""
+        try:
+            print("📥 Loading transformer model...")
+            self.model = SymmetricTransformer2DModel.from_pretrained(
+                self.transformer_path,
+                subfolder="transformer",
+            )
+            print("📥 Loading VQ model...")
+            self.vq_model = VQModel.from_pretrained(
+                self.model_path, 
+                subfolder="vqvae"
+            )
+            print("📥 Loading text encoder...")
+            self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
+                self.model_path, 
+                subfolder="text_encoder"
+            )
+            print("📥 Loading tokenizer...")
+            self.tokenizer = CLIPTokenizer.from_pretrained(
+                self.model_path, 
+                subfolder="tokenizer"
+            )
+            print("📥 Loading scheduler...")
+            self.scheduler = Scheduler.from_pretrained(
+                self.model_path, 
+                subfolder="scheduler"
+            )
+
+            print("🔧 Assembling pipeline...")
+            self.pipe = UnifiedPipeline(
+                vqvae=self.vq_model,
+                tokenizer=self.tokenizer,
+                text_encoder=self.text_encoder,
+                transformer=self.model,
+                scheduler=self.scheduler,
+            )
+            print(f"🚀 Moving models to {self.device}...")
+            self.pipe.to(self.device)
+        except Exception as e:
+            print(f"❌ Error loading models: {str(e)}")
+            raise
+    
+    def text_to_image(self, prompt, negative_prompt, height, width, steps, cfg_scale, seed):
+        """Generate image from text prompt"""
+        try:
+            if seed == -1:
+                generator = None
+            else:
+                generator = torch.manual_seed(seed)
+            
+            if not negative_prompt:
+                negative_prompt = "worst quality, low quality, low res, blurry, distortion, watermark, logo, signature, text, jpeg artifacts, signature, sketch, duplicate, ugly, identifying mark"
+            
+            output = self.pipe(
+                prompt=[prompt],
+                negative_prompt=negative_prompt,
+                height=height,
+                width=width,
+                guidance_scale=cfg_scale,
+                num_inference_steps=steps,
+                mask_token_embedding=None,
+                generator=generator
+            )
+            
+            if hasattr(output, 'images') and len(output.images) > 0:
+                return output.images[0]
+            else:
+                return None
+                
+        except Exception as e:
+            gr.Error(f"Error generating image: {str(e)}")
+            return None
+    
+    def image_to_text(self, image, question, height, width, steps, cfg_scale):
+        """Answer question about the image"""
+        try:
+            if image is None:
+                return "Please upload an image."
+            
+            # Convert PIL image to tensor
+            if isinstance(image, np.ndarray):
+                image = Image.fromarray(image)
+            
+            # Save image temporarily and load using the existing function
+            temp_path = "temp_image.jpg"
+            image.save(temp_path)
+            
+            try:
+                images = load_images_to_tensor(temp_path, target_size=(height, width))
+            finally:
+                if os.path.exists(temp_path):
+                    os.remove(temp_path)
+            
+            if images is None:
+                return "Failed to process the image."
+            
+            questions = [question] * len(images)
+            
+            output = self.pipe(
+                prompt=questions,
+                image=images,
+                height=height,
+                width=width,
+                guidance_scale=cfg_scale,
+                num_inference_steps=steps,
+                mask_token_embedding=None,
+            )
+            
+            if hasattr(output, 'prompts') and len(output.prompts) > 0:
+                return output.prompts[0]
+            else:
+                return "No response generated."
+                
+        except Exception as e:
+            return f"Error processing image: {str(e)}"
+
+
+def create_muddit_interface():
+    # Initialize the model interface
+    interface = MudditInterface()
+    
+    with gr.Blocks(title="Muddit Interface", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🎨 Muddit Interface")
+        gr.Markdown("Generate images from text or ask questions about images using Muddit.")
+        
+        with gr.Tabs():
+            # Text-to-Image Tab
+            with gr.TabItem("🖼️ Text-to-Image"):
+                gr.Markdown("### Generate images from text descriptions")
+                
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        t2i_prompt = gr.Textbox(
+                            label="Prompt",
+                            placeholder="A majestic night sky awash with billowing clouds, sparkling with a million twinkling stars",
+                            lines=3
+                        )
+                        t2i_negative = gr.Textbox(
+                            label="Negative Prompt (optional)",
+                            placeholder="worst quality, low quality, blurry...",
+                            lines=2
+                        )
+                        
+                        with gr.Row():
+                            t2i_width = gr.Slider(
+                                minimum=256, maximum=1024, value=1024, step=64,
+                                label="Width"
+                            )
+                            t2i_height = gr.Slider(
+                                minimum=256, maximum=1024, value=1024, step=64,
+                                label="Height"
+                            )
+                        
+                        with gr.Row():
+                            t2i_steps = gr.Slider(
+                                minimum=1, maximum=100, value=64, step=1,
+                                label="Inference Steps"
+                            )
+                            t2i_cfg = gr.Slider(
+                                minimum=1.0, maximum=20.0, value=9.0, step=0.5,
+                                label="CFG Scale"
+                            )
+                        
+                        t2i_seed = gr.Number(
+                            label="Seed (-1 for random)", 
+                            value=42, 
+                            precision=0
+                        )
+                        
+                        t2i_generate = gr.Button("🎨 Generate Image", variant="primary")
+                    
+                    with gr.Column(scale=1):
+                        t2i_output = gr.Image(label="Generated Image", type="pil")
+                
+                t2i_generate.click(
+                    fn=interface.text_to_image,
+                    inputs=[t2i_prompt, t2i_negative, t2i_height, t2i_width, t2i_steps, t2i_cfg, t2i_seed],
+                    outputs=[t2i_output]
+                )
+            
+            # Visual Question Answering Tab
+            with gr.TabItem("❓ Visual Question Answering"):
+                gr.Markdown("### Ask questions about images")
+                
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        vqa_image = gr.Image(
+                            label="Upload Image", 
+                            type="pil"
+                        )
+                        vqa_question = gr.Textbox(
+                            label="Question",
+                            placeholder="What do you see in this image?",
+                            lines=2
+                        )
+                        
+                        with gr.Row():
+                            vqa_width = gr.Slider(
+                                minimum=256, maximum=1024, value=1024, step=64,
+                                label="Width"
+                            )
+                            vqa_height = gr.Slider(
+                                minimum=256, maximum=1024, value=1024, step=64,
+                                label="Height"
+                            )
+                        
+                        with gr.Row():
+                            vqa_steps = gr.Slider(
+                                minimum=1, maximum=100, value=64, step=1,
+                                label="Inference Steps"
+                            )
+                            vqa_cfg = gr.Slider(
+                                minimum=1.0, maximum=20.0, value=9.0, step=0.5,
+                                label="CFG Scale"
+                            )
+                        
+                        vqa_submit = gr.Button("🤔 Ask Question", variant="primary")
+                    
+                    with gr.Column(scale=1):
+                        vqa_output = gr.Textbox(
+                            label="Answer", 
+                            lines=5,
+                            interactive=False
+                        )
+                
+                vqa_submit.click(
+                    fn=interface.image_to_text,
+                    inputs=[vqa_image, vqa_question, vqa_height, vqa_width, vqa_steps, vqa_cfg],
+                    outputs=[vqa_output]
+                )
+        
+        # Example section
+        with gr.Accordion("📝 Examples", open=False):
+            gr.Markdown("""
+            ### Text-to-Image Examples:
+            - "A majestic night sky awash with billowing clouds, sparkling with a million twinkling stars"
+            - "A hyper realistic image of a chimpanzee with a glass-enclosed brain on his head"
+            - "A samurai in a stylized cyberpunk outfit adorned with intricate steampunk gear"
+            
+            ### VQA Examples:
+            - "What objects do you see in this image?"
+            - "How many people are in the picture?"
+            - "What is the main subject of this image?"
+            - "Describe the scene in detail"
+            """)
+    
+    return demo
+
+
+if __name__ == "__main__":
+    demo = create_muddit_interface()
+    demo.launch() 

requirements.txt

@@ -0,0 +1,9 @@
+gradio>=4.0.0
+torch>=2.0.0
+torchvision>=0.15.0
+transformers>=4.40.0
+diffusers>=0.30.0
+pillow>=9.0.0
+numpy>=1.21.0
+accelerate>=0.20.0
+safetensors>=0.3.0 

src/pipeline.py

@@ -0,0 +1,1231 @@
+# 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,
+    CLIPImageProcessor,
+    CLIPVisionModelWithProjection,
+)
+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("is such is", "").replace("such is", "")
+    # 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: GemmaTokenizerFast
+    text_encoder: CLIPTextModelWithProjection
+    text_encoder_2: Gemma2Model
+    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: GemmaTokenizerFast = None,
+        text_encoder_2: Gemma2Model = 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)
+                
+                text_encoder_2_mask_id = self.tokenizer_2.convert_tokens_to_ids("<extra_id_0>")
+                self.scheduler.config.mask_token_id = text_encoder_2_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:
+                    try:
+                        if mask_token_embedding.endswith(".pth"):
+                            mask_token_embedding = torch.load(mask_token_embedding)
+                        else:
+                            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)
+
+                    except Exception as e:
+                        print(f"Error loading mask token embedding: {e}")
+                        print("Using random initialized mask token embedding")
+                        mask_token_embedding = None
+
+                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)
+
+
+class UnifiedPipeline_new(DiffusionPipeline):
+    image_processor: VaeImageProcessor
+    vqvae: VQModel
+    tokenizer: CLIPTokenizer
+    tokenizer_2: GemmaTokenizerFast
+    text_encoder: CLIPTextModelWithProjection
+    text_encoder_2: Gemma2Model
+    image_encoder: CLIPVisionModelWithProjection
+    clip_image_processor: CLIPImageProcessor
+    transformer: SymmetricTransformer2DModel
+    scheduler: Scheduler
+
+    def __init__(
+        self,
+        vqvae: VQModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModelWithProjection,
+        transformer: SymmetricTransformer2DModel,
+        scheduler: Scheduler,
+        tokenizer_2: Optional[GemmaTokenizerFast]=None,
+        text_encoder_2: Optional[Gemma2Model]=None,
+        image_encoder: Optional[CLIPVisionModelWithProjection]=None,
+        clip_image_processor: Optional[CLIPImageProcessor]=None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vqvae=vqvae,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            image_encoder=image_encoder,
+            clip_image_processor=clip_image_processor,
+            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[torch.Tensor] = 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 = "<mask>" * 256
+                prompt = [prompt] * image.shape[0]
+                
+                text_encoder_2_mask_id = self.tokenizer_2.convert_tokens_to_ids("<mask>")
+                self.scheduler.config.mask_token_id = text_encoder_2_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:
+                    try:
+                        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)
+
+                    except Exception as e:
+                        print(f"Error loading mask token embedding: {e}")
+                        print("Using random initialized mask token embedding")
+                        mask_token_embedding = None
+
+                self.scheduler.config.mask_token_id = clip_mask_id
+                
+                input_ids = torch.ones(
+                    size=(image.shape[0], 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] * image.shape[0]
+                elif isinstance(prompt, str):
+                    question_ids = torch.LongTensor([self.tokenizer.encode(prompt)])
+                    question_ids = question_ids.repeat(image.shape[0], 1)
+                    
+                    q_len = len(question_ids[0]) - 1   # remove <eos> token
+                    question_len = [q_len] * image.shape[0]
+                    
+                    input_ids[:, :q_len] = question_ids[:, :-1]
+                else: 
+                    assert isinstance(prompt, list), f"prompt must be None or str or list!"
+                    assert len(prompt) == image.shape[0], 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 image is not None:
+            batch_size = image.shape[0]
+        elif isinstance(prompt, list):
+            batch_size = len(prompt)
+        elif isinstance(prompt, str):
+            batch_size = 1
+            prompt = [prompt]
+        else:
+            raise ValueError("prompt must be None or str or list!")
+
+        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 and text2image:
+            if text_encoder_type == "open_clip":
+                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":
+                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=self.transformer.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)
+                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)
+
+                        prompt_embeds = self.image_encoder(
+                            self.clip_image_processor(
+                                image,
+                                return_tensors="pt",
+                                do_rescale=False,
+                                do_resize=True,
+                                do_normalize=True,
+                            ).to(self._execution_device, dtype=self.image_encoder.dtype).pixel_values
+                        ).image_embeds  # [b, 1024]
+                    
+                    if text_encoder_type in ("t5_clip", "gemma"):
+                        outputs = self.text_encoder_2(latents, return_dict=True)
+                        encoder_hidden_states = outputs.last_hidden_state
+                    else:
+                        outputs = self.text_encoder(latents, return_dict=True, output_hidden_states=True)
+                        encoder_hidden_states = outputs.hidden_states[-2]
+                else:
+                    model_input = latents
+                    
+                img_ids = _prepare_latent_image_ids(
+                    model_input.shape[0], 
+                    model_input.shape[-2],
+                    model_input.shape[-1],
+                    self._execution_device,
+                    self.transformer.dtype
+                )
+                txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(
+                    device=self._execution_device,
+                    dtype=self.transformer.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], device=self._execution_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: Tuple[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.num_attention_heads * self.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.pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(self.joint_attention_dim, self.inner_dim)
+     
+        self.transformer_blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                )
+                for i in range(self.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.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.codebook_size
+        )
+        self.cond_embed = TimestepEmbedding( 
+            micro_cond_embed_dim + self.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.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.num_attention_heads * self.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.inner_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.joint_attention_dim, bias=False)
+            )
+        else:
+            self.adapter = None
+
+        self.context_embedder = nn.Linear(self.joint_attention_dim, self.inner_dim)
+     
+        self.transformer_blocks = nn.ModuleList(
+            [
+                TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.num_attention_heads,
+                    attention_head_dim=self.attention_head_dim,
+                )
+                for i in range(self.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.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.codebook_size
+        )
+        self.cond_embed = TimestepEmbedding( 
+            micro_cond_embed_dim + self.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.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]

train/dataset_utils.py

@@ -0,0 +1,472 @@
+# 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 torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+from PIL.ImageOps import exif_transpose
+from PIL import Image
+import io
+import json
+import numpy as np
+import pyarrow.parquet as pq
+import random
+import bisect
+import pyarrow.fs as fs
+
+
+@torch.no_grad()
+def tokenize_prompt(
+    tokenizer, 
+    prompt, 
+    text_encoder_architecture='open_clip',
+    padding='max_length',
+    max_length=77,
+    max_length_t5=256,
+):
+    if text_encoder_architecture == 'CLIP' or text_encoder_architecture == 'open_clip':
+        input_ids = tokenizer(
+            prompt,
+            truncation=True,
+            padding=padding,
+            max_length=max_length,
+            return_tensors="pt",
+        ).input_ids
+        return input_ids
+    elif text_encoder_architecture == 't5_clip': # we have two tokenizers, 1st for CLIP, 2nd for T5
+        input_ids = []
+        input_ids.append(tokenizer[0](
+            prompt,
+            truncation=True,
+            padding=padding,
+            max_length=max_length,
+            return_tensors="pt",
+        ).input_ids)
+        input_ids.append(tokenizer[1](
+            prompt,
+            truncation=True,
+            padding=padding,
+            max_length=max_length_t5,
+            return_tensors="pt",
+        ).input_ids)
+        return input_ids
+    elif text_encoder_architecture == "gemma":
+        input_ids = []
+        input_ids.append(tokenizer[0](
+            prompt,
+            truncation=True,
+            padding=padding,
+            max_length=max_length,
+            return_tensors="pt",
+        ).input_ids)
+        input_ids.append(tokenizer[1](
+            prompt,
+            truncation=True,
+            padding=padding,
+            max_length=max_length_t5,
+            return_tensors="pt",
+        ).input_ids)
+        return input_ids
+    else:
+        raise ValueError(f"Unknown text_encoder_architecture: {text_encoder_architecture}")
+
+
+def encode_prompt(
+    text_encoder, 
+    input_ids, 
+    text_encoder_architecture='open_clip'
+):
+    if text_encoder_architecture == 'CLIP' or text_encoder_architecture == 'open_clip':
+        outputs = text_encoder(input_ids=input_ids, return_dict=True, output_hidden_states=True)
+        encoder_hidden_states = outputs.hidden_states[-2]
+        cond_embeds = outputs[0]
+        return encoder_hidden_states, cond_embeds
+    elif text_encoder_architecture == 't5_clip':
+        outputs_clip = text_encoder[0](
+            input_ids=input_ids[0], 
+            return_dict=True, 
+            output_hidden_states=True
+        )
+        outputs_t5 = text_encoder[1](
+            input_ids=input_ids[1], 
+            return_dict=True, 
+            output_hidden_states=True
+        )
+        encoder_hidden_states = outputs_t5.last_hidden_state
+        cond_embeds = outputs_clip.text_embeds
+        return encoder_hidden_states, cond_embeds
+    elif text_encoder_architecture == "gemma":
+        outputs_clip = text_encoder[0](
+            input_ids=input_ids[0], 
+            return_dict=True, 
+            output_hidden_states=True
+        )
+        outputs_gemma = text_encoder[1](
+            input_ids=input_ids[1], 
+            return_dict=True, 
+            output_hidden_states=True
+        )
+        encoder_hidden_states = outputs_gemma.last_hidden_state
+        cond_embeds = outputs_clip.text_embeds
+        return encoder_hidden_states, cond_embeds
+    else:
+        raise ValueError(f"Unknown text_encoder_architecture: {text_encoder_architecture}")
+
+
+def process_image(image, size, Norm=False, hps_score=6.0): 
+    image = exif_transpose(image)
+
+    if not image.mode == "RGB":
+        image = image.convert("RGB")
+
+    orig_height = image.height
+    orig_width = image.width
+
+    image = transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR)(image)
+
+    c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(size, size))
+    image = transforms.functional.crop(image, c_top, c_left, size, size)
+    image = transforms.ToTensor()(image)
+
+    if Norm:
+        image = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)(image)
+
+    micro_conds = torch.tensor(
+        [orig_width, orig_height, c_top, c_left, hps_score],
+    )
+
+    return {"image": image, "micro_conds": micro_conds}    
+    
+
+class ImageCaptionLargeDataset(Dataset):
+    def __init__(
+        self, 
+        root_dir,
+        tokenizer,
+        size,
+        text_encoder_architecture="CLIP",
+        norm=False
+    ):
+        self.root_dir = root_dir
+        self.tokenizer = tokenizer
+        self.size = size
+        self.text_encoder_architecture = text_encoder_architecture
+        self.norm = norm
+        
+        self.data_list = []        
+        for root, dirnames, filenames in os.walk(root_dir):
+            for filename in filenames:
+                if filename.endswith(".jpg") or filename.endswith(".png"):
+                    base_name = os.path.splitext(filename)[0]
+                    txt_file = os.path.join(root, base_name + ".txt")
+                    if os.path.exists(txt_file):
+                        self.data_list.append((root, base_name + ".txt", filename))
+            
+    def __len__(self):
+        return len(self.data_list)
+    
+    def __getitem__(self, idx):
+        try:
+            sub_dir, txtfilename, imgfilename = self.data_list[idx]
+            img_path = os.path.join(sub_dir, imgfilename)
+            caption_path = os.path.join(sub_dir, txtfilename)
+            
+            image = Image.open(img_path).convert("RGB")
+            ret = process_image(image, self.size, self.norm)
+            
+            with open(caption_path, "r", encoding="utf-8") as f:
+                caption = f.read().strip()
+            
+            ret["prompt_input_ids"] = tokenize_prompt(
+                self.tokenizer, caption, self.text_encoder_architecture
+            )
+            
+            return ret
+        
+        except Exception as e:
+            print("===========================================")
+            print(f"[Warning] Error at index {idx}: {img_path}")
+            print("===========================================")
+            if idx + 1 < len(self.data_list):
+                return self.__getitem__(idx + 1)
+            else:
+                return self.__getitem__(len(self.data_list) - 1)
+
+
+class MultiSourceVLDataset(Dataset):
+    """
+    A unified dataloader for
+      • LLaVA-Instruct-150K
+      • MMMU (multiple-choice QA)
+      • VQAv2
+      • Local caption files under `pdd3/`
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        size: int,
+        text_encoder_architecture: str = "CLIP",
+        norm: bool = False,
+        # ----- paths -----
+        llava_json: str = None, llava_img_root: str = None,
+        mmmu_json: str = None,  mmmu_img_root: str = None,
+        vqa_ann_json: str = None, vqa_img_root: str = None,
+        gqa_json: str = None, gqa_img_root: str = None,
+        coco_json: str = None, coco_img_root: str = None,
+        coco_qa_json: str = None,
+        mg_llava_json: str = None, mg_llava_root: str = None,
+        pdd3_dir: str = None, caption_dir: str = None,
+    ):
+        self.tokenizer = tokenizer
+        self.size = size
+        self.arch = text_encoder_architecture
+        self.norm = norm
+
+        self.gen_samples = []      # [(img_path, prompt), ...]
+        self.mmu_samples = []      # [(img_path, question, answer), ...]
+
+        if llava_json:
+            self._load_llava(llava_json, llava_img_root)
+        if mmmu_json:
+            self._load_mmmu(mmmu_json, mmmu_img_root)
+        if vqa_ann_json:
+            self._load_vqav2(vqa_ann_json, vqa_img_root)
+        if coco_json:
+            self._load_coco2014_captions(coco_json, coco_img_root)
+        if coco_qa_json:
+            self._load_coco2014_qa(coco_qa_json, coco_img_root)
+        if gqa_json:
+            self._load_gqa(gqa_json, gqa_img_root)
+        if mg_llava_json:
+            self._load_mg_llava(mg_llava_json, mg_llava_root)
+        if caption_dir:
+            self._load_caption(caption_dir)
+        if pdd3_dir:
+            self._load_pdd3(pdd3_dir)
+
+        self.len_mmu = len(self.mmu_samples)
+        self.len_gen = len(self.gen_samples)
+
+    # ------------------------------------------------------------------ #
+    #                          dataset parsers                            #
+    # ------------------------------------------------------------------ #
+    def _load_llava(self, json_path, img_root):
+        with open(json_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            
+        for ex in data:
+            img_file = os.path.join(img_root, ex["image"])
+
+            human_msg = next(m["value"] for m in ex["conversations"] if m["from"] == "human")
+            gpt_msg   = next(m["value"] for m in ex["conversations"] if m["from"] == "gpt")
+
+            self.mmu_samples.append((img_file, human_msg.strip(), gpt_msg.strip())) 
+
+    def _load_mmmu(self, json_path, img_root):
+        with open(json_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            
+        for ex in data:
+            img_file = os.path.join(img_root, ex["image"])
+            choices  = "\n".join([f"{chr(65+i)}. {c}" for i, c in enumerate(ex["choices"])])
+            
+            question = f"{ex['question'].strip()}\n{choices}"
+            answer = f"{ex['answer']}"
+            
+            self.mmu_samples.append((img_file, question, answer))
+
+    def _load_coco2014_qa(self, ann_jsonl, img_root):
+        with open(ann_jsonl, "r", encoding="utf-8") as file:
+            data = [json.loads(line) for line in file if line.strip()]
+        
+        for ann in data:
+            image = ann["image"]
+            question = ann["question"]
+            answer = ann["label"]
+            
+            image_path = os.path.join(img_root, image)
+            self.mmu_samples.append((image_path, question, answer))
+            
+
+    def _load_coco2014_captions(self, ann_json, img_root):
+        """
+        Load COCO 2014 image-caption pairs from caption annotation file.
+
+        Args:
+            ann_json (str): Path to COCO-style captions JSON (e.g., captions_train2014.json)
+            img_root (str): Directory containing COCO images (should include 'train2014/' and 'val2014/' subdirs)
+        """
+        with open(ann_json, "r") as f:
+            data = json.load(f)
+
+        is_train = "train" in os.path.basename(ann_json).lower()
+        img_subdir = "train2014" if is_train else "val2014"
+        prefix = "COCO_train2014_" if is_train else "COCO_val2014_"
+
+        for ann in data["annotations"]:
+            image_id = ann["image_id"]
+            caption = ann["caption"]
+            
+            image_filename = f"{prefix}{image_id:012d}.jpg"
+            image_path = os.path.join(img_root, img_subdir, image_filename)
+
+            question = "Please describe this image concisely."
+            self.mmu_samples.append((image_path, question, caption))
+
+    def _load_vqav2(self, ann_json, img_root):
+        with open(ann_json, "r") as file:
+            annos = json.load(file)
+
+        for ann in annos:
+            q = ann["question"]
+            answer = ann["answer"]
+            img_path = ann["image"]
+            img_file = os.path.join(
+                img_root,
+                img_path   # if val, modify to val2014
+            )
+
+            self.mmu_samples.append((img_file, q, answer))
+
+    def _load_gqa(self, ann_json_root, img_root):
+        annos = {}
+        
+        for jsonfile in os.listdir(ann_json_root):
+            jsonpath = os.path.join(ann_json_root, jsonfile)
+            with open(jsonpath, "r") as file:
+                anno = json.load(file)
+                annos.update(anno)
+        
+        for ann in annos.values():
+            q = ann["question"]
+            answer = ann["fullAnswer"]
+            img_name = ann["imageId"] + ".jpg"
+            img_path = os.path.join(
+                img_root, 
+                img_name
+            )
+            
+            self.mmu_samples.append((img_path, q, answer))
+
+    def _load_mg_llava(self, json_path, img_root):
+        with open(json_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+            
+        for ex in data:
+            image = ex.get("image", None)
+            if image is not None:
+                img_file = os.path.join(img_root, ex["image"])
+                if os.path.exists(img_file):
+                    human_msg = next(m["value"] for m in ex["conversations"] if m["from"] == "human")
+                    gpt_msg   = next(m["value"] for m in ex["conversations"] if m["from"] == "gpt")
+
+                    self.mmu_samples.append((img_file, human_msg.strip(), gpt_msg.strip())) 
+
+    def _load_caption(self, root_dir):
+        for root, _, files in os.walk(root_dir):
+            for f in files:
+                if f.lower().endswith((".jpg", ".png")):
+                    base = os.path.splitext(f)[0]
+                    txt_path = os.path.join(root, base + ".txt")
+                    if os.path.exists(txt_path):
+                        with open(txt_path, "r") as file:
+                            caption = file.read().strip()
+                        q = "Please describe this image."
+                        self.mmu_samples.append((os.path.join(root, f), q, caption))
+
+    def _load_pdd3(self, root_dir):
+        for root, _, files in os.walk(root_dir):
+            for f in files:
+                if f.lower().endswith((".jpg", ".png")):
+                    base = os.path.splitext(f)[0]
+                    txt_path = os.path.join(root, base + ".txt")
+                    if os.path.exists(txt_path):
+                        with open(txt_path, "r") as file:
+                            caption = file.read().strip()
+                        self.gen_samples.append((os.path.join(root, f), caption))
+
+    # ------------------------------------------------------------------ #
+    #                       PyTorch Dataset API                          #
+    # ------------------------------------------------------------------ #
+    def __len__(self):
+        return max(self.len_gen, self.len_mmu)
+
+    def __getitem__(self, idx):
+        get_mmu_data = False
+        get_gen_data = False
+        
+        while not get_mmu_data:
+            try:
+                mmu_img_path, question, answer = self.mmu_samples[idx]
+                get_mmu_data = True
+            except:
+                idx = random.randint(0, self.len_mmu - 1)
+            
+        while not get_gen_data:
+            try:
+                gen_img_path, prompt = self.gen_samples[idx]
+                get_gen_data = True
+            except:
+                idx = random.randint(0, self.len_gen - 1)
+
+        try:
+            # ---- image ----
+            mmu_image = Image.open(mmu_img_path).convert("RGB")
+            mmu_ret = process_image(mmu_image, self.size, self.norm)
+            
+            gen_image = Image.open(gen_img_path).convert("RGB")
+            gen_ret = process_image(gen_image, self.size, self.norm)
+
+            ret = dict(
+                gen_image=gen_ret["image"],
+                gen_micro_conds=gen_ret["micro_conds"],
+                mmu_image=mmu_ret["image"],
+                mmu_micro_conds=mmu_ret["micro_conds"]
+            )
+
+            # ---- text ----
+            question = question.replace("<image>", "").replace("\n", "")
+            question_ids = tokenize_prompt(
+                self.tokenizer, 
+                question, 
+                self.arch,
+                padding=False,
+            )
+            question_ids = question_ids[:, :-1]
+            q_len = len(question_ids[0])
+            if answer:
+                full_prompt = question + " " + answer
+            else:
+                full_prompt = question
+            mmu_input_ids = tokenize_prompt(self.tokenizer, full_prompt, self.arch)
+            
+            gen_input_ids = tokenize_prompt(self.tokenizer, prompt, self.arch)
+
+            ret.update({
+                "gen_input_ids": gen_input_ids,
+                "mmu_input_ids": mmu_input_ids,
+                "question_len": torch.LongTensor([q_len])
+            })
+            return ret
+        except:
+            print("================================================================")
+            print(f"There is something wrong with {mmu_img_path} or {gen_img_path}.")
+            print("================================================================")
+            if idx < self.len_gen - 1 or idx < self.len_mmu - 1:
+                return self.__getitem__(idx + 1)
+            else:
+                idx = random.randint(0, self.len_gen - 1)
+                return self.__getitem__(idx)

train/instruction_tuning.py

@@ -0,0 +1,1246 @@
+# 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 argparse
+import copy
+import logging
+import math
+import os
+from pathlib import Path
+import sys
+sys.path.append(os.getcwd())
+import json
+import gc
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from peft import LoraConfig
+from peft.utils import get_peft_model_state_dict
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5Tokenizer,
+)
+
+import diffusers.optimization
+from diffusers import VQModel
+
+from src.scheduler import Scheduler
+from diffusers.loaders import LoraLoaderMixin
+from diffusers.utils import is_wandb_available
+from src.pipeline import UnifiedPipeline
+from torchvision.utils import save_image, make_grid
+from train.trainer_utils import save_checkpoint
+from train.dataset_utils import ImageCaptionLargeDataset, MultiSourceVLDataset
+from train.dataset_utils import tokenize_prompt, encode_prompt
+from src.transformer import SymmetricTransformer2DModel
+from train.trainer_utils import load_images_to_tensor
+
+if is_wandb_available():
+    import wandb
+    # wandb.login(key="")
+
+logger = get_logger(__name__, log_level="INFO")
+
+import torch._dynamo
+torch._dynamo.config.verbose = True
+
+# Optionally suppress errors to fall back to eager execution
+torch._dynamo.config.suppress_errors = True
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--pretrained_transformer_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained transformer.",
+    )
+    parser.add_argument(
+        "--text_encoder_architecture",
+        type=str,
+        default="open_clip",
+        required=False,
+        help="The architecture of the text encoder. One of ['CLIP', 'open_clip', 'flan-t5-base','Qwen2-0.5B','gemini-2b',long_t5_clip','t5_clip']",
+    )
+    parser.add_argument(
+        "--dataset_type",
+        type=str,
+        default=None,
+        required=False,
+        help="The type of the dataset.",
+    )
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--caption_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--llava_json_path",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--llava_image_root",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--mmmu_json_path",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--mmmu_image_root",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--vqa_ann_json_path",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--vqa_image_root",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--coco_json",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--coco_qa_json",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--coco_img_root",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--gqa_json_root",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--gqa_image_root",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--mg_llava_json",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--mg_llava_root",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--training_from_scratch",
+        type=bool,
+        default=False,
+        required=False
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--variant",
+        type=str,
+        default=None,
+        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.9999)
+    parser.add_argument("--ema_update_after_step", type=int, default=0)
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="muse_training",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
+            "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
+            "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
+            "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
+            "instructions."
+        ),
+    )
+    parser.add_argument(
+        "--logging_steps",
+        type=int,
+        default=50,
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more details"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--text_loss_weight",
+        type=float,
+        default=0.2,
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=0.0003,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--validation_steps",
+        type=int,
+        default=100,
+        help=(
+            "Run validation every X steps. Validation consists of running the prompt"
+            " `args.validation_prompt` multiple times: `args.num_validation_images`"
+            " and logging the images."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="wandb",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument("--validation_prompts", type=str, nargs="*")
+    parser.add_argument("--validation_vqa_prompts", type=str, default=None)
+    parser.add_argument("--validation_images", type=str, default=None)
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument("--split_vae_encode", type=int, required=False, default=None)
+    parser.add_argument("--min_masking_rate", type=float, default=0.0)
+    parser.add_argument("--cond_dropout_prob", type=float, default=0.0)
+    parser.add_argument("--max_grad_norm", default=50.0, type=float, help="Max gradient norm.", required=False)
+    parser.add_argument("--use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--text_encoder_use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--lora_r", default=16, type=int)
+    parser.add_argument("--lora_alpha", default=32, type=int)
+    parser.add_argument("--lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--text_encoder_lora_r", default=16, type=int)
+    parser.add_argument("--text_encoder_lora_alpha", default=32, type=int)
+    parser.add_argument("--text_encoder_lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--train_text_encoder", action="store_true")
+    parser.add_argument("--image_to_text_only", action="store_true")
+    parser.add_argument("--image_key", type=str, required=False)
+    parser.add_argument("--prompt_key", type=str, required=False)
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--prompt_prefix", type=str, required=False, default=None)
+
+    args = parser.parse_args()
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+    if args.instance_data_dir is not None:
+        if not os.path.exists(args.instance_data_dir):
+            raise ValueError(f"Does not exist: `--args.instance_data_dir` {args.instance_data_dir}")
+
+    return args
+
+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 main(args):
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if accelerator.is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers("meissonic", config=vars(copy.deepcopy(args)))
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    if args.text_encoder_architecture == "open_clip":
+        text_encoder = CLIPTextModelWithProjection.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
+        )
+        tokenizer = CLIPTokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
+        )
+        tokenizer_2 = None
+        text_encoder_2 = None
+        
+        text_encoder.requires_grad_(False)
+    elif args.text_encoder_architecture == "t5_clip":
+        tokenizer = CLIPTokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
+        )
+        text_encoder = CLIPTextModelWithProjection.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
+        )   
+        
+        tokenizer_2 = T5Tokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer_2", variant=args.variant,
+        )
+        text_encoder_2 = T5EncoderModel.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder_2", variant=args.variant,
+        )
+        
+        text_encoder.requires_grad_(False)
+        text_encoder_2.requires_grad_(False)
+    else:
+        raise ValueError(f"Unknown text encoder architecture: {args.text_encoder_architecture}")
+    
+    vq_model = VQModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vqvae", revision=args.revision, variant=args.variant
+    )
+    vq_model.requires_grad_(False)
+
+    model = SymmetricTransformer2DModel.from_pretrained(
+        args.pretrained_transformer_path, 
+        subfolder="transformer", 
+        low_cpu_mem_usage=False, 
+        device_map=None
+    )
+    
+    if model.config.tokenizer_vocab_size is None:
+        if args.text_encoder_architecture == "open_clip":
+            model.register_to_config(tokenizer_vocab_size=len(tokenizer))
+            # model.config.tokenizer_vocab_size = len(tokenizer)    # We exclude the mask token in the predicted logits
+        elif args.text_encoder_architecture == "t5_clip":
+            model.register_to_config(tokenizer_vocab_size=len(tokenizer_2))
+            # model.config.tokenizer_vocab_size = len(tokenizer_2)      # We don't need to add new token
+            if model.adapter is None:
+                raise ValueError(f"The MMDiT must has adapter if you want to use t5_clip mode!!!")
+        else:
+            raise ValueError(f"Unknown text encoder architecture!")
+    
+        print(f"model's tokenizer vocab size is {model.config.tokenizer_vocab_size}")
+        model.text_decoder = nn.Sequential(
+            nn.LayerNorm(model.inner_dim, elementwise_affine=False, eps=1e-6),
+            nn.Linear(model.inner_dim, model.config.tokenizer_vocab_size, bias=False)
+        )
+
+    model = torch.compile(model)
+
+    if args.use_lora:
+        lora_config = LoraConfig(
+            r=args.lora_r,
+            lora_alpha=args.lora_alpha,
+            target_modules=args.lora_target_modules,
+        )
+        model.add_adapter(lora_config)
+
+    model.train()
+
+    if args.image_to_text_only:
+        frozen_keys = ["project_from_hidden", "up_block", "mlm_layer"]
+        for n, p in model.named_parameters():
+            if any([frozen_key in n for frozen_key in frozen_keys]):
+                p.requires_grad_(False)
+            else:
+                p.requires_grad_(True)
+    else:
+        model.requires_grad_(True)
+
+    if args.gradient_checkpointing:
+        model.enable_gradient_checkpointing()   
+
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            transformer_lora_layers_to_save = None
+            text_encoder_lora_layers_to_save = None
+
+            for model_ in models:
+                if isinstance(model_, type(accelerator.unwrap_model(model))):
+                    if args.use_lora:
+                        transformer_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "transformer"))
+                elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                    if args.text_encoder_use_lora:
+                        text_encoder_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "text_encoder"))
+                else:
+                    raise ValueError(f"unexpected save model: {model_.__class__}")
+
+                # make sure to pop weight so that corresponding model is not saved again
+                weights.pop()
+
+            if transformer_lora_layers_to_save is not None or text_encoder_lora_layers_to_save is not None:
+                LoraLoaderMixin.save_lora_weights(
+                    output_dir,
+                    unet_lora_layers=transformer_lora_layers_to_save,
+                    text_encoder_lora_layers=text_encoder_lora_layers_to_save,
+                )
+                
+
+    def load_model_hook(models, input_dir):
+        transformer = None
+        text_encoder_ = None
+
+        # this part is added for keep consistency when add model.compile() in the model
+        def adap_compile(ori_dict):#add '_orig_mod.' to each key
+            new_dict = {}
+            for k,v in ori_dict.items():
+                new_dict['_orig_mod.' + k] = v
+            return new_dict
+            
+        while len(models) > 0:
+            model_ = models.pop()
+
+            if isinstance(model_, type(accelerator.unwrap_model(model))):
+                if args.use_lora:
+                    transformer = model_
+                else:
+                    load_model = SymmetricTransformer2DModel.from_pretrained(os.path.join(input_dir, "transformer"), low_cpu_mem_usage=False, device_map=None)
+                    model_.load_state_dict(adap_compile(load_model.state_dict()))
+                    del load_model
+            elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                if args.text_encoder_use_lora:
+                    text_encoder_ = model_
+                else:
+                    try:
+                        load_model = CLIPTextModelWithProjection.from_pretrained(os.path.join(input_dir, "text_encoder"))
+                        model_.load_state_dict(load_model.state_dict())
+                        # print('finished loading text encoder!')
+                    except:
+                        print('Not found text-encoder model in current folder. So we download one text encoder from Internet.')
+                        load_model = CLIPTextModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
+                        model_.load_state_dict(load_model.state_dict())
+                    del load_model
+            else:
+                raise ValueError(f"unexpected save model: {model.__class__}")
+
+        if transformer is not None or text_encoder_ is not None:
+            lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir)
+            LoraLoaderMixin.load_lora_into_text_encoder(
+                lora_state_dict, network_alphas=network_alphas, text_encoder=text_encoder_
+            )
+            LoraLoaderMixin.load_lora_into_transformer(
+                lora_state_dict, network_alphas=network_alphas, transformer=transformer
+            )
+
+    accelerator.register_load_state_pre_hook(load_model_hook)
+    accelerator.register_save_state_pre_hook(save_model_hook)
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        )
+
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+            )
+
+        optimizer_cls = bnb.optim.AdamW8bit
+    else:
+        optimizer_cls = torch.optim.AdamW
+
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for p in model.parameters() if p.requires_grad],
+            "weight_decay": args.adam_weight_decay,
+        }
+    ]
+    optimizer = optimizer_cls(
+        optimizer_grouped_parameters,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    logger.info("Creating dataloaders and lr_scheduler")
+
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    if args.text_encoder_architecture == "t5_clip":
+        tokenizer_for_dataset = [tokenizer, tokenizer_2]
+    else:
+        tokenizer_for_dataset = tokenizer
+
+    if args.dataset_type == "ImageCaptionLargeDataset":
+        dataset = ImageCaptionLargeDataset(
+            root_dir=args.instance_data_dir,
+            tokenizer=tokenizer_for_dataset,
+            size=args.resolution,
+            text_encoder_architecture=args.text_encoder_architecture
+        )
+    elif args.dataset_type == "MultiSourceVLDataset":
+        dataset = MultiSourceVLDataset(
+            tokenizer=tokenizer_for_dataset,
+            size=args.resolution,
+            text_encoder_architecture=args.text_encoder_architecture,
+            norm=False,
+            llava_json=args.llava_json_path,
+            llava_img_root=args.llava_image_root,
+            mmmu_json=args.mmmu_json_path,
+            mmmu_img_root=args.mmmu_image_root,
+            vqa_ann_json=args.vqa_ann_json_path,
+            vqa_img_root=args.vqa_image_root,
+            coco_json=args.coco_json,
+            coco_qa_json=args.coco_qa_json,
+            coco_img_root=args.coco_img_root,
+            gqa_json=args.gqa_json_root,
+            gqa_img_root=args.gqa_image_root,
+            mg_llava_json=args.mg_llava_json,
+            mg_llava_root=args.mg_llava_root,
+            caption_dir=args.caption_dir,
+            pdd3_dir=args.instance_data_dir,
+        )
+    elif args.dataset_type == "DATA_TYPE":
+        raise NotImplementedError("DATA_TYPE is not yet supported")
+    else:
+        assert False
+
+    def collate_fn(samples):
+        gen_images = [sample["gen_image"] for sample in samples]
+        mmu_images = [sample["mmu_image"] for sample in samples]
+        
+        gen_micro_conds = [sample["gen_micro_conds"] for sample in samples]
+        mmu_micro_conds = [sample["mmu_micro_conds"] for sample in samples]
+        
+        gen_images = torch.stack(gen_images, dim=0)
+        mmu_images = torch.stack(mmu_images, dim=0)
+        
+        gen_micro_conds = torch.stack(gen_micro_conds, dim=0)
+        mmu_micro_conds = torch.stack(mmu_micro_conds, dim=0)
+        
+        if isinstance(samples[0]["gen_input_ids"], list):
+            gen_input_ids = [sample["gen_input_ids"][0] for sample in samples]
+            gen_input_ids_2 = [sample["gen_input_ids"][1] for sample in samples]
+            
+            gen_input_ids = torch.cat(gen_input_ids, dim=0)
+            gen_input_ids_2 = torch.cat(gen_input_ids_2, dim=0)
+            gen_input_ids = [gen_input_ids, gen_input_ids_2]
+        else:
+            gen_input_ids = [sample["gen_input_ids"] for sample in samples]
+            mmu_input_ids = [sample["mmu_input_ids"] for sample in samples]
+            
+            gen_input_ids = torch.cat(gen_input_ids, dim=0)
+            mmu_input_ids = torch.cat(mmu_input_ids, dim=0)
+        
+        if samples[0].get("question_len", None) is not None:
+            question_len = [sample["question_len"] for sample in samples]
+            
+            question_len = torch.cat(question_len, dim=0)   # [B, ]
+        else:
+            question_len = None
+        
+        ret = dict(
+            gen_images=gen_images,
+            mmu_images=mmu_images,
+            gen_micro_conds=gen_micro_conds,
+            mmu_micro_conds=mmu_micro_conds,
+            gen_input_ids=gen_input_ids,
+            mmu_input_ids=mmu_input_ids,
+            question_len=question_len
+        )
+        
+        return ret
+    
+    train_dataloader = DataLoader(
+        dataset,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=collate_fn,
+        pin_memory=True,
+    )
+    train_dataloader.num_batches = len(train_dataloader)
+
+    lr_scheduler = diffusers.optimization.get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+    )
+
+    logger.info("Preparing model, optimizer and dataloaders")
+
+    model, optimizer, lr_scheduler, train_dataloader = accelerator.prepare(
+        model, optimizer, lr_scheduler, train_dataloader
+    )
+
+    train_dataloader.num_batches = len(train_dataloader)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    if args.text_encoder_architecture == "t5_clip":
+        text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+        text_encoder_2.to(device=accelerator.device, dtype=weight_dtype)
+    else:
+        text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+            
+    vq_model.to(device=accelerator.device)
+
+    with torch.no_grad():
+        if args.text_encoder_architecture == "t5_clip":
+            _input_ids_tmp_ = tokenize_prompt([tokenizer, tokenizer_2], "", args.text_encoder_architecture)
+            _input_ids_tmp_[0] = _input_ids_tmp_[0].to(accelerator.device)
+            _input_ids_tmp_[1] = _input_ids_tmp_[1].to(accelerator.device)
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                [text_encoder, text_encoder_2], 
+                _input_ids_tmp_, 
+                args.text_encoder_architecture
+            )
+        else:
+            _input_ids_tmp_ = tokenize_prompt(tokenizer, "", args.text_encoder_architecture)
+            _input_ids_tmp_ = _input_ids_tmp_.to(accelerator.device)
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                text_encoder,
+                _input_ids_tmp_,
+                args.text_encoder_architecture
+            )
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(train_dataloader.num_batches / args.gradient_accumulation_steps)
+    # Afterwards we recalculate our number of training epochs.
+    # Note: We are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info(f"  Num training steps = {args.max_train_steps}")
+    logger.info(f"  Instantaneous batch size per device = { args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+
+    resume_from_checkpoint = args.resume_from_checkpoint
+    if resume_from_checkpoint:
+        if resume_from_checkpoint == "latest":
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            if len(dirs) > 0:
+                resume_from_checkpoint = os.path.join(args.output_dir, dirs[-1])
+            else:
+                resume_from_checkpoint = None
+
+        if resume_from_checkpoint is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+        else:
+            accelerator.print(f"Resuming from checkpoint {resume_from_checkpoint}")
+
+    if resume_from_checkpoint is None:
+        global_step = 0
+        first_epoch = 0
+    else:
+        accelerator.load_state(resume_from_checkpoint)
+        global_step = int(os.path.basename(resume_from_checkpoint).split("-")[1])
+        first_epoch = global_step // num_update_steps_per_epoch
+
+    # As stated above, we are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    for epoch in range(first_epoch, num_train_epochs):
+        for batch in train_dataloader:
+            torch.cuda.empty_cache()
+            with torch.no_grad():
+                gen_pixel_values = batch["gen_images"].to(accelerator.device, non_blocking=True)
+                mmu_pixel_values = batch["mmu_images"].to(accelerator.device, non_blocking=True)
+                
+                gen_micro_conds = batch["gen_micro_conds"].to(accelerator.device, non_blocking=True)
+                mmu_micro_conds = batch["mmu_micro_conds"].to(accelerator.device, non_blocking=True)
+    
+                # ====================== tokenize images ======================
+                pixel_values = torch.cat([gen_pixel_values, mmu_pixel_values], dim=0)
+                batch_size = pixel_values.shape[0]
+                
+                split_batch_size = args.split_vae_encode if args.split_vae_encode is not None else batch_size
+                num_splits = math.ceil(batch_size / split_batch_size)
+                image_tokens = []
+                for i in range(num_splits):
+                    start_idx = i * split_batch_size
+                    end_idx = min((i + 1) * split_batch_size, batch_size)
+                    image_tokens.append(
+                        vq_model.quantize(
+                            vq_model.encode(pixel_values[start_idx:end_idx]).latents
+                        )[2][2].reshape(split_batch_size, -1)
+                    )
+                image_tokens = torch.cat(image_tokens, dim=0)
+                gen_image_tokens, mmu_image_tokens = image_tokens.chunk(2, dim=0)
+                # ====================== tokenize images ======================
+
+
+                # ====================== encode clean text prompts ======================
+                if args.text_encoder_architecture == "t5_clip":
+                    gen_input_ids_clip = batch["gen_input_ids"][0].to(accelerator.device, non_blocking=True)
+                    gen_input_ids_t5 = batch["gen_input_ids"][1].to(accelerator.device, non_blocking=True)
+                    encoder_hidden_states, cond_embeds = encode_prompt(
+                        [text_encoder, text_encoder_2], 
+                        [gen_input_ids_clip, gen_input_ids_t5], 
+                        args.text_encoder_architecture
+                    )
+                else:
+                    gen_input_ids = batch["gen_input_ids"].to(accelerator.device, non_blocking=True)          
+                    gen_encoder_hidden_states, gen_cond_embeds = encode_prompt(
+                        text_encoder, 
+                        gen_input_ids, 
+                        args.text_encoder_architecture
+                    )
+                gen_encoder_hidden_states = gen_encoder_hidden_states.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                gen_cond_embeds = gen_cond_embeds.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                # ====================== encode clean text prompts ======================
+
+
+                # ====================== image perturbation ======================
+                half_batch_size, seq_len = gen_image_tokens.shape
+                sigma = torch.rand(half_batch_size, device=gen_image_tokens.device)
+                image_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                image_mask_prob = image_mask_prob.clip(args.min_masking_rate)
+
+                num_token_masked = (seq_len * image_mask_prob).round().clamp(min=1)
+                batch_randperm = torch.rand(half_batch_size, seq_len, device=gen_image_tokens.device).argsort(dim=-1)
+                mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                mask_id = accelerator.unwrap_model(model).config.vocab_size - 1
+                gen_image_ids = torch.where(mask, mask_id, gen_image_tokens)
+                image_labels = torch.where(mask, gen_image_tokens, -100)
+                # ====================== image perturbation ======================
+
+
+                # ====================== text perturbation ======================        
+                if args.text_encoder_architecture == "t5_clip":
+                    mmu_input_ids_clip = batch["mmu_input_ids"][0].to(accelerator.device, non_blocking=True)
+                    mmu_input_ids_t5 = batch["mmu_input_ids"][1].to(accelerator.device, non_blocking=True)
+                    half_batch_size, seq_len = mmu_input_ids_t5.shape
+                    sigma = torch.rand(half_batch_size, device=mmu_image_tokens.device)
+                    text_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                    text_mask_prob = text_mask_prob.clip(args.min_masking_rate)
+                    text_timestep = text_mask_prob.clone().clamp(min=1e-3)
+
+                    num_token_masked = (seq_len * text_mask_prob).round().clamp(min=1)
+                    batch_randperm = torch.rand(half_batch_size, seq_len, device=mmu_image_tokens.device).argsort(dim=-1)
+                    mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                    extra_id_0_token = "<extra_id_0>"
+                    t5_mask_id = tokenizer_2.convert_tokens_to_ids(extra_id_0_token)
+                    masked_prompt_input_ids_t5 = torch.where(mask, t5_mask_id, mmu_input_ids_t5)
+                    text_labels = torch.where(mask, mmu_input_ids_t5, -100)
+                    
+                    # prepare input_ids for clip model
+                    batch_prompt_2 = []
+                    for i in range(masked_prompt_input_ids_t5.size(0)):
+                        masked_prompt_input_id = masked_prompt_input_ids_t5[i].tolist()
+                        prompt_2 = tokenizer_2.decode(masked_prompt_input_id, skip_special_tokens=True)
+                        batch_prompt_2.append(prompt_2)
+                    
+                    masked_prompt_input_ids_clip = tokenizer(
+                        batch_prompt_2,
+                        truncation=True,
+                        padding="max_length",
+                        max_length=77,
+                        return_tensors="pt"
+                    ).input_ids
+                    masked_prompt_input_ids_clip = masked_prompt_input_ids_clip.to(accelerator.device)
+                else:
+                    extra_id_0_token = "<extra_id_0>"
+                    num_new_tokens = tokenizer.add_tokens(extra_id_0_token)
+                    clip_mask_id = tokenizer.convert_tokens_to_ids(extra_id_0_token)
+                    if num_new_tokens > 0:
+                        text_encoder.resize_token_embeddings(len(tokenizer))
+                        mask_token_embedding = text_encoder.get_input_embeddings().weight[clip_mask_id]
+                        mask_token_embedding = mask_token_embedding.clone().detach().cpu().float()
+                        if accelerator.is_main_process:
+                            print("Saving masked token embedding...")
+                            torch.save(mask_token_embedding, os.path.join(args.output_dir, "mask_token_embedding.pth"))
+                            
+
+                    mmu_input_ids = batch["mmu_input_ids"].to(accelerator.device, non_blocking=True) # [B, L]
+                    question_len = batch["question_len"]    # [B, ]
+                    if question_len is not None:
+                        question_len = question_len.to(accelerator.device, non_blocking=True)
+                                            
+                    half_batch_size, seq_len = mmu_input_ids.shape
+                    answer_len = seq_len - question_len
+                    
+                    sigma = torch.rand(half_batch_size, device=mmu_image_tokens.device)
+                    text_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                    text_mask_prob = text_mask_prob.clip(args.min_masking_rate)
+                    text_timestep = text_mask_prob.clone().clamp(min=1e-3)
+                    
+                    num_token_masked = ((seq_len - question_len) * text_mask_prob).round().clamp(min=1) # [B, ]
+                    num_token_masked = torch.minimum(num_token_masked, answer_len)
+
+                    seq_idx = torch.arange(seq_len, device=mmu_image_tokens.device).unsqueeze(0).repeat(half_batch_size, 1)
+                    answer_region = seq_idx >= question_len.unsqueeze(1)
+                    
+                    rand_value = torch.rand(half_batch_size, seq_len, device=mmu_image_tokens.device)
+                    rand_value = rand_value.masked_fill(~answer_region, float("inf"))
+                    
+                    order = rand_value.argsort(dim=-1)
+                    order = order.argsort(dim=-1)
+                    mask = order < num_token_masked.unsqueeze(-1)
+                     
+                    # mask = torch.zeros_like(mmu_input_ids)
+                    # for b in range(half_batch_size):
+                    #     ans_len = seq_len - question_len[b]
+                    #     batch_randperm = torch.rand(1, ans_len, device=mmu_image_tokens.device).argsort(dim=-1)
+                    #     mask[b, question_len[b]:] = batch_randperm < num_token_masked[b].unsqueeze(-1)
+                            
+                    mmu_input_ids_clip = torch.where(mask, clip_mask_id, mmu_input_ids)
+                    text_labels = torch.where(mask, mmu_input_ids, -100)
+                # ====================== text perturbation ======================
+
+
+                # ====================== encode masked text prompts ======================
+                if args.text_encoder_architecture == "t5_clip":
+                    masked_encoder_hidden_states, masked_cond_embeds = encode_prompt(
+                        [text_encoder, text_encoder_2], 
+                        [masked_prompt_input_ids_clip, masked_prompt_input_ids_t5], 
+                        args.text_encoder_architecture
+                    )
+                else:
+                    mmu_encoder_hidden_states, mmu_cond_embeds = encode_prompt(
+                        text_encoder, 
+                        mmu_input_ids_clip, 
+                        args.text_encoder_architecture
+                    )
+                mmu_encoder_hidden_states = mmu_encoder_hidden_states.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                mmu_cond_embeds = mmu_cond_embeds.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                # ====================== encode masked text prompts ======================
+                
+                
+                # for CFG
+                if args.cond_dropout_prob > 0.0:
+                    assert encoder_hidden_states is not None
+
+                    batch_size = encoder_hidden_states.shape[0]
+
+                    mask = (
+                        torch.zeros((batch_size, 1, 1), device=encoder_hidden_states.device).float().uniform_(0, 1)
+                        < args.cond_dropout_prob
+                    )
+
+                    empty_embeds_ = empty_embeds.expand(batch_size, -1, -1)
+                    encoder_hidden_states = torch.where(
+                        (encoder_hidden_states * mask).bool(), encoder_hidden_states, empty_embeds_
+                    )
+
+                    empty_clip_embeds_ = empty_clip_embeds.expand(batch_size, -1)
+                    cond_embeds = torch.where((cond_embeds * mask.squeeze(-1)).bool(), cond_embeds, empty_clip_embeds_)
+                    
+                vae_scale_factor = 2 ** (len(vq_model.config.block_out_channels) - 1)
+                resolution = args.resolution // vae_scale_factor
+                gen_image_ids = gen_image_ids.reshape(half_batch_size, resolution, resolution)
+                mmu_image_ids = mmu_image_tokens.reshape(half_batch_size, resolution, resolution)
+
+
+            # Train Step
+            with accelerator.accumulate(model):
+                codebook_size = accelerator.unwrap_model(model).config.codebook_size                   
+                if args.resolution == 1024: # only stage 3 and stage 4 do not apply 2*
+                    img_ids = _prepare_latent_image_ids(
+                        gen_image_ids.shape[0], 
+                        gen_image_ids.shape[-2],
+                        gen_image_ids.shape[-1],
+                        gen_image_ids.device,
+                        gen_image_ids.dtype
+                    )
+                else:
+                    img_ids = _prepare_latent_image_ids(
+                        gen_image_ids.shape[0],
+                        gen_image_ids.shape[-2],
+                        gen_image_ids.shape[-1],
+                        gen_image_ids.device,
+                        gen_image_ids.dtype
+                    )
+
+                txt_ids = torch.zeros(gen_encoder_hidden_states.shape[1], 3).to(device=gen_image_ids.device, dtype=gen_image_ids.dtype)
+                
+                image_logits = model(
+                    hidden_states=gen_image_ids, # should be (batch size, channel, height, width)
+                    encoder_hidden_states=gen_encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                    micro_conds=gen_micro_conds,
+                    pooled_projections=gen_cond_embeds, # should be (batch_size, projection_dim)
+                    img_ids=img_ids,
+                    txt_ids=txt_ids,
+                    timestep=image_mask_prob * 1000,
+                )[0]
+                image_logits = image_logits.reshape(half_batch_size, codebook_size, -1)
+                image_logits = image_logits.permute(0, 2, 1)
+                image_logits = image_logits.reshape(-1, codebook_size)
+
+                image_loss = F.cross_entropy(
+                    image_logits,
+                    image_labels.view(-1),
+                    ignore_index=-100,
+                    reduction="mean",
+                )
+
+                text_logits = model(
+                    hidden_states=mmu_image_ids, # should be (batch size, channel, height, width)
+                    encoder_hidden_states=mmu_encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                    micro_conds=mmu_micro_conds,
+                    pooled_projections=mmu_cond_embeds, # should be (batch_size, projection_dim)
+                    img_ids=img_ids,
+                    txt_ids=txt_ids,
+                    timestep=text_timestep * 1000,
+                )[1]
+                text_logits = text_logits.reshape(-1, accelerator.unwrap_model(model).config.tokenizer_vocab_size)
+
+                text_loss = F.cross_entropy(
+                    text_logits,
+                    text_labels.view(-1),
+                    ignore_index=-100,
+                    reduction="none",
+                )
+                text_loss = text_loss.reshape(half_batch_size, -1).mean(-1)
+                text_loss = text_loss / text_timestep
+                text_loss = text_loss.mean()
+
+                loss = image_loss + args.text_loss_weight * text_loss
+                
+                # Gather the losses across all processes for logging (if we use distributed training).
+                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                avg_masking_rate = accelerator.gather(text_mask_prob.repeat(args.train_batch_size)).mean()
+
+                accelerator.backward(loss)
+                
+                # Temporarily add this to identify unused parameters
+                # for name, param in accelerator.unwrap_model(model).named_parameters():
+                #     if param.grad is None:
+                #         print(f"Unused parameter: {name}")
+
+                if args.max_grad_norm is not None and accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                if (global_step + 1) % args.logging_steps == 0:
+                    logs = {
+                        "step_loss": avg_loss.item(),
+                        "lr": lr_scheduler.get_last_lr()[0],
+                        "avg_masking_rate": avg_masking_rate.item(),
+                    }
+                    accelerator.log(logs, step=global_step + 1)
+
+                    logger.info(
+                        f"Step: {global_step + 1} "
+                        f"Loss: {avg_loss.item():0.4f} "
+                        f"LR: {lr_scheduler.get_last_lr()[0]:0.6f}"
+                    )
+
+                if (global_step + 1) % args.checkpointing_steps == 0:
+                    save_checkpoint(args, accelerator, global_step + 1, logger)
+
+                if (global_step + 1) % args.validation_steps == 0 and accelerator.is_main_process:
+
+                    with torch.no_grad():
+                        logger.info("Generating images...")
+
+                        model.eval()
+
+                        scheduler = Scheduler.from_pretrained(
+                            args.pretrained_model_name_or_path,
+                            subfolder="scheduler",
+                            revision=args.revision,
+                            variant=args.variant,
+                        )
+                        
+                        pipe = UnifiedPipeline(
+                            transformer=accelerator.unwrap_model(model),
+                            tokenizer=tokenizer,
+                            text_encoder=text_encoder,
+                            vqvae=vq_model,
+                            scheduler=scheduler,
+                            tokenizer_2=tokenizer_2,
+                            text_encoder_2=text_encoder_2,
+                        )
+
+                        if not args.image_to_text_only:
+                            output = pipe(
+                                prompt=args.validation_prompts,
+                                height=args.resolution,
+                                width=args.resolution,
+                                guidance_scale=9,
+                                num_inference_steps=64,
+                            )
+                            pil_images = output.images
+
+                            result=[]
+                            for img in pil_images:
+                                if not isinstance(img, torch.Tensor):
+                                    img = transforms.ToTensor()(img)
+                                result.append(img.unsqueeze(0))
+                            result = torch.cat(result,dim=0)
+                            result = make_grid(result, nrow=3)
+                            save_image(result,os.path.join(args.output_dir, str(global_step)+'_text2image_1024_CFG-9.png'))
+                        
+                            output_data = {
+                                "step": global_step,
+                                "prompts": args.validation_prompts,
+                                "images": [f"{global_step}_text2image_1024_CFG-9_{i}.png" for i in range(len(pil_images))]
+                            }
+
+                            with open(os.path.join(args.output_dir, f"text2image_{global_step}.json"), "w") as f:
+                                json.dump(output_data, f, indent=2)
+
+                        image = load_images_to_tensor(args.validation_images, target_size=(args.resolution, args.resolution))        
+                        output = pipe(
+                            prompt=args.validation_vqa_prompts,
+                            height=args.resolution, 
+                            width=args.resolution, 
+                            guidance_scale=9,
+                            image=image,
+                            num_inference_steps=64
+                        )
+                        prompts = output.prompts
+                        
+                        output_data = {
+                            "step": global_step,
+                            "prompts": prompts,
+                        }
+
+                        with open(os.path.join(args.output_dir, f"image2text_{global_step}.json"), "w") as f:
+                            json.dump(output_data, f, indent=2)
+
+                        model.train()
+
+                global_step += 1
+
+            # Stop training if max steps is reached
+            if global_step >= args.max_train_steps:
+                break
+        # End for
+
+    accelerator.wait_for_everyone()
+
+    # Evaluate and save checkpoint at the end of training
+    save_checkpoint(args, accelerator, global_step, logger)
+
+    # Save the final trained checkpoint
+    if accelerator.is_main_process:
+        model = accelerator.unwrap_model(model)
+        model.save_pretrained(args.output_dir)
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main(parse_args())

train/instruction_tuning.sh

@@ -0,0 +1,37 @@
+# bash it in root path
+PYTHON_PATH='./' accelerate launch --multi_gpu --gpu_ids '0,1,2,3,4,5,6,7,8' --main_process_port 25000 --num_processes 8 train/instruction_tuning.py \
+        --output_dir "/path/to/output/dir" \
+        --train_batch_size 8 \
+        --gradient_accumulation_steps 8 \
+        --learning_rate 1e-4 \
+        --text_loss_weight 0.2 \
+        --max_grad_norm 10 \
+        --pretrained_model_name_or_path "MeissonFlow/Meissonic" \
+        --pretrained_transformer_path "MeissonFlow/Meissonic" \
+        --text_encoder_architecture 'open_clip' \
+        --dataset_type 'MultiSourceVLDataset' \
+        --instance_data_dir '/path/to/data' \
+        --llava_json_path '/path/to/llava_instruct_150k.json' \
+        --llava_image_root '/path/to/coco/train2017' \
+        --resolution 512 \
+        --mixed_precision fp16 \
+        --lr_scheduler constant \
+        --use_8bit_adam \
+        --dataloader_num_workers 4 \
+        --validation_prompts \
+            'a boy' \
+            'A serene mountain landscape with towering snow-capped peaks, a crystal-clear blue lake reflecting the mountains, dense pine forests, and a vibrant orange sunrise illuminating the sky.' \
+            'A playful golden retriever puppy with a shiny coat, bounding through a meadow filled with colorful wildflowers, under a bright, clear blue sky.' \
+            'A bustling city street at night, illuminated by vibrant neon signs in various colors, with busy pedestrians, street vendors, and a light rain creating reflective puddles on the pavement.' \
+            'A majestic, medieval castle perched on a rugged cliffside, overlooking a vast, calm ocean at sunset, with the sky painted in hues of pink, orange, and purple.' \
+            'An elegant ballerina in a white tutu, dancing gracefully on a grand stage with ornate, gold-trimmed curtains, under a spotlight that casts a soft glow.' \
+            'A cozy, rustic log cabin nestled in a snow-covered forest, with smoke rising from the stone chimney, warm lights glowing from the windows, and a path of footprints leading to the front door.'\
+            'A Cute Cat' \
+            'A Snow Mountain'\
+        --validation_images '/path/to/validation/images/dir' \
+        --validation_vqa_prompts 'Please describe this image.' \
+        --max_train_steps 100000 \
+        --checkpointing_steps 100 \
+        --validation_steps 100 \
+        --report_to 'wandb' \
+        --logging_steps 10

train/train_text_decoder.py

@@ -0,0 +1,1017 @@
+# 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 argparse
+import copy
+import logging
+import math
+import os
+import json
+from contextlib import nullcontext
+from pathlib import Path
+import sys
+sys.path.append(os.getcwd())
+import gc
+import torch
+import torch.nn.functional as F
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from peft import LoraConfig
+from peft.utils import get_peft_model_state_dict
+from torch.utils.data import DataLoader, default_collate
+from torchvision import transforms
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+)
+import diffusers.optimization
+from diffusers import EMAModel, VQModel 
+from src.scheduler import Scheduler
+from diffusers.loaders import LoraLoaderMixin
+from diffusers.utils import is_wandb_available
+from src.pipeline import UnifiedPipeline
+from torchvision.utils import save_image,make_grid
+from datasets import load_dataset
+from train.trainer_utils import save_checkpoint
+from train.dataset_utils import ImageCaptionDataset, HuggingFaceDataset
+from train.dataset_utils import tokenize_prompt, encode_prompt
+from src.transformer import SymmetricTransformer2DModel, Transformer2DModel
+
+if is_wandb_available():
+    import wandb
+    # wandb.login(key="")
+
+logger = get_logger(__name__, log_level="INFO")
+
+import torch._dynamo
+torch._dynamo.config.verbose = True
+
+# Optionally suppress errors to fall back to eager execution
+torch._dynamo.config.suppress_errors = True
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_architecture",
+        type=str,
+        default="Meissonic",
+        required=False
+    )
+    parser.add_argument(
+        "--text_encoder_architecture",
+        type=str,
+        default="open_clip",
+        required=False,
+        help="The architecture of the text encoder. One of ['CLIP', 'open_clip', 'flan-t5-base','Qwen2-0.5B','gemini-2b',long_CLIP_T5_base','CLIP_T5_base']",
+    )
+    parser.add_argument(
+        "--instance_dataset",
+        type=str,
+        default=None,
+        required=False,
+        help="The dataset to use for training. One of ['MSCOCO600K', 'PickaPicV2']",
+    )
+    parser.add_argument(
+        "--training_from_scratch",
+        type=bool,
+        default=False,
+        required=False
+    )
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--variant",
+        type=str,
+        default=None,
+        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+    )
+    parser.add_argument(
+        "--instance_data_dataset",
+        type=str,
+        default=None,
+        required=False,
+        help="A Hugging Face dataset containing the training images",
+    )
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--instance_data_image", type=str, default=None, required=False, help="A single training image"
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.9999)
+    parser.add_argument("--ema_update_after_step", type=int, default=0)
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="muse_training",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
+            "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
+            "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
+            "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
+            "instructions."
+        ),
+    )
+    parser.add_argument(
+        "--logging_steps",
+        type=int,
+        default=50,
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more details"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=0.0003,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--validation_steps",
+        type=int,
+        default=100,
+        help=(
+            "Run validation every X steps. Validation consists of running the prompt"
+            " `args.validation_prompt` multiple times: `args.num_validation_images`"
+            " and logging the images."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="wandb",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument("--validation_prompts", type=str, nargs="*")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument("--split_vae_encode", type=int, required=False, default=None)
+    parser.add_argument("--min_masking_rate", type=float, default=0.0)
+    parser.add_argument("--cond_dropout_prob", type=float, default=0.0)
+    parser.add_argument("--max_grad_norm", default=50.0, type=float, help="Max gradient norm.", required=False)
+    parser.add_argument("--use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--text_encoder_use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--lora_r", default=16, type=int)
+    parser.add_argument("--lora_alpha", default=32, type=int)
+    parser.add_argument("--lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--text_encoder_lora_r", default=16, type=int)
+    parser.add_argument("--text_encoder_lora_alpha", default=32, type=int)
+    parser.add_argument("--text_encoder_lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--train_text_encoder", action="store_true")
+    parser.add_argument("--image_key", type=str, required=False)
+    parser.add_argument("--prompt_key", type=str, required=False)
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--prompt_prefix", type=str, required=False, default=None)
+
+    args = parser.parse_args()
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+    num_datasources = sum(
+        [x is not None for x in [args.instance_data_dir, args.instance_data_image, args.instance_data_dataset]]
+    )
+
+    if num_datasources != 1:
+        raise ValueError(
+            "provide one and only one of `--instance_data_dir`, `--instance_data_image`, or `--instance_data_dataset`"
+        )
+
+    if args.instance_data_dir is not None:
+        if not os.path.exists(args.instance_data_dir):
+            raise ValueError(f"Does not exist: `--args.instance_data_dir` {args.instance_data_dir}")
+
+    if args.instance_data_image is not None:
+        if not os.path.exists(args.instance_data_image):
+            raise ValueError(f"Does not exist: `--args.instance_data_image` {args.instance_data_image}")
+
+    if args.instance_data_dataset is not None and (args.image_key is None or args.prompt_key is None):
+        raise ValueError("`--instance_data_dataset` requires setting `--image_key` and `--prompt_key`")
+
+    return args
+
+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
+    )
+    # latent_image_ids = latent_image_ids.unsqueeze(0).repeat(batch_size, 1, 1)
+
+    return latent_image_ids.to(device=device, dtype=dtype)
+
+def main(args):
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    # if args.pretrained_model_architecture == "Meissonic":
+    #     from src.pipeline import Pipeline
+    # else:
+    #     raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")
+
+
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if accelerator.is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers("meissonic", config=vars(copy.deepcopy(args)))
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    if args.text_encoder_architecture == "open_clip":
+        if args.resume_from_checkpoint:
+            text_encoder = CLIPTextModelWithProjection.from_pretrained( 
+                args.resume_from_checkpoint, subfolder="text_encoder", variant=args.variant
+            )
+            tokenizer = CLIPTokenizer.from_pretrained(
+                args.resume_from_checkpoint, subfolder="tokenizer", variant=args.variant
+            )
+        else:
+            text_encoder = CLIPTextModelWithProjection.from_pretrained(
+                args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
+            )
+            tokenizer = CLIPTokenizer.from_pretrained(
+                args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
+            )
+    else:
+        raise ValueError(f"Unknown text encoder architecture: {args.text_encoder_architecture}")
+    
+    vq_model = VQModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vqvae", revision=args.revision, variant=args.variant
+    )
+
+    if args.train_text_encoder:
+        if args.text_encoder_use_lora:
+            lora_config = LoraConfig(
+                r=args.text_encoder_lora_r,
+                lora_alpha=args.text_encoder_lora_alpha,
+                target_modules=args.text_encoder_lora_target_modules,
+            )
+            text_encoder.add_adapter(lora_config)
+        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+            text_encoder[0].train()
+            text_encoder[0].requires_grad_(True)
+            text_encoder[1].train()
+            text_encoder[1].requires_grad_(True)
+        else:
+            text_encoder.train()
+            text_encoder.requires_grad_(True)
+    else:
+        if args.text_encoder_architecture == "CLIP_T5_base":  # Not support yet. Only support open_clip
+            text_encoder[0].eval()
+            text_encoder[0].requires_grad_(False)
+            text_encoder[1].eval()
+            text_encoder[1].requires_grad_(False)
+        else:
+            text_encoder.eval()
+            text_encoder.requires_grad_(False)
+
+    vq_model.requires_grad_(False)
+
+    if args.pretrained_model_architecture == "Meissonic":
+        if args.training_from_scratch:
+            model = SymmetricTransformer2DModel( 
+                patch_size = 1,
+                in_channels = 64,
+                num_layers = 14, 
+                num_single_layers = 28,
+                attention_head_dim = 128, 
+                num_attention_heads = 8, 
+                joint_attention_dim = 1024,
+                pooled_projection_dim = 1024,
+                guidance_embeds = False,
+                axes_dims_rope = (16, 56, 56),
+                downsample= True,
+                upsample= True,
+            )
+        else:
+            orig_model = Transformer2DModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="transformer", low_cpu_mem_usage=False, device_map=None)
+            orig_config = orig_model.config
+            config = {}
+            for k, v in orig_config.items():
+                if k.startswith("_"):
+                    continue
+                config[k] = v
+            config["tokenizer_vocab_size"] = tokenizer.vocab_size
+
+            model = SymmetricTransformer2DModel(**config)
+            model.load_state_dict(orig_model.state_dict(), strict=False)
+
+            del orig_model
+            gc.collect()
+            torch.cuda.empty_cache()
+    else:
+        raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")
+
+    model = torch.compile(model)
+
+    if args.use_lora:
+        lora_config = LoraConfig(
+            r=args.lora_r,
+            lora_alpha=args.lora_alpha,
+            target_modules=args.lora_target_modules,
+        )
+        model.add_adapter(lora_config)
+
+    model.train()
+    for n, p in model.named_parameters():
+        if "text_decoder" in n:
+            p.requires_grad_(True)
+        else:
+            p.requires_grad_(False)
+
+    if args.gradient_checkpointing:
+        model.enable_gradient_checkpointing()
+        if args.train_text_encoder:
+            if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+                text_encoder[0].gradient_checkpointing_enable()
+                text_encoder[1].gradient_checkpointing_enable()
+            else:
+                text_encoder.gradient_checkpointing_enable()
+
+    if args.use_ema: # Not verify the robostness of this part
+        ema = EMAModel(
+            model.parameters(),
+            decay=args.ema_decay,
+            update_after_step=args.ema_update_after_step,
+            model_cls= Transformer2DModel, 
+            model_config=model.config,
+        )
+
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            transformer_lora_layers_to_save = None
+            text_encoder_lora_layers_to_save = None
+
+            for model_ in models:
+                if isinstance(model_, type(accelerator.unwrap_model(model))):
+                    if args.use_lora:
+                        transformer_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "transformer"))
+                elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                    if args.text_encoder_use_lora:
+                        text_encoder_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "text_encoder"))
+                else:
+                    raise ValueError(f"unexpected save model: {model_.__class__}")
+
+                # make sure to pop weight so that corresponding model is not saved again
+                weights.pop()
+
+            if transformer_lora_layers_to_save is not None or text_encoder_lora_layers_to_save is not None:
+                LoraLoaderMixin.save_lora_weights(
+                    output_dir,
+                    unet_lora_layers=transformer_lora_layers_to_save,
+                    text_encoder_lora_layers=text_encoder_lora_layers_to_save,
+                )
+
+            if args.use_ema:
+                ema.save_pretrained(os.path.join(output_dir, "ema_model"))
+
+    def load_model_hook(models, input_dir):
+        transformer = None
+        text_encoder_ = None
+
+        # this part is added for keep consistency when add model.compile() in the model
+        def adap_compile(ori_dict):#add '_orig_mod.' to each key
+            new_dict = {}
+            for k,v in ori_dict.items():
+                new_dict['_orig_mod.'+k] = v
+            return new_dict
+            
+        while len(models) > 0:
+            model_ = models.pop()
+
+            if isinstance(model_, type(accelerator.unwrap_model(model))):
+                if args.use_lora:
+                    transformer = model_
+                else:
+                    if args.pretrained_model_architecture == "Meissonic":
+                        load_model = SymmetricTransformer2DModel.from_pretrained(os.path.join(input_dir, "transformer"), low_cpu_mem_usage=False, device_map=None)
+                    else:
+                        raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")
+                    model_.load_state_dict(adap_compile(load_model.state_dict()))
+                    del load_model
+            elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                if args.text_encoder_use_lora:
+                    text_encoder_ = model_
+                else:
+                    try:
+                        load_model = CLIPTextModelWithProjection.from_pretrained(os.path.join(input_dir, "text_encoder"))
+                        model_.load_state_dict(load_model.state_dict())
+                        # print('finished loading text encoder!')
+                    except:
+                        print('Not found text-encoder model in current folder. So we download one text encoder from Internet.')
+                        load_model = CLIPTextModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
+                        model_.load_state_dict(load_model.state_dict())
+                    del load_model
+            else:
+                raise ValueError(f"unexpected save model: {model.__class__}")
+
+        if transformer is not None or text_encoder_ is not None:
+            lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir)
+            LoraLoaderMixin.load_lora_into_text_encoder(
+                lora_state_dict, network_alphas=network_alphas, text_encoder=text_encoder_
+            )
+            LoraLoaderMixin.load_lora_into_transformer(
+                lora_state_dict, network_alphas=network_alphas, transformer=transformer
+            )
+
+        if args.use_ema:
+            load_from = EMAModel.from_pretrained(os.path.join(input_dir, "ema_model"), model_cls=Transformer2DModel)
+            ema.load_state_dict(adap_compile(load_from.state_dict()))
+            del load_from
+
+    accelerator.register_load_state_pre_hook(load_model_hook)
+    accelerator.register_save_state_pre_hook(save_model_hook)
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        )
+
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+            )
+
+        optimizer_cls = bnb.optim.AdamW8bit
+    else:
+        optimizer_cls = torch.optim.AdamW
+
+
+    optimizer_parameters = [p for p in model.parameters() if p.requires_grad]
+
+    optimizer = optimizer_cls(
+        optimizer_parameters,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    logger.info("Creating dataloaders and lr_scheduler")
+
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    if args.instance_dataset == "MyParquetDataset":
+        dataset = ImageCaptionDataset(
+            root_dir=args.instance_data_dir,  # something like '../parquets_father_dir/'
+            tokenizer=tokenizer,
+            size=args.resolution,
+            text_encoder_architecture=args.text_encoder_architecture
+        )
+    elif args.instance_dataset == 'HuggingFaceDataset':  # you can try this first, just download dataset from huggingface
+        dataset = HuggingFaceDataset(
+            hf_dataset=load_dataset(args.instance_data_dir, split="train"),  # something like './parquets_father_dir/'
+            tokenizer=tokenizer,
+            image_key='image',
+            prompt_key='caption',
+            prompt_prefix=args.prompt_prefix,
+            size=args.resolution,
+            text_encoder_architecture=args.text_encoder_architecture
+        )
+    elif args.instance_dataset == "DATA_TYPE":
+        raise NotImplementedError("DATA_TYPE is not yet supported")
+    else:
+        assert False
+
+    train_dataloader = DataLoader(
+        dataset,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=default_collate,
+        pin_memory=True,
+    )
+    train_dataloader.num_batches = len(train_dataloader)
+
+    lr_scheduler = diffusers.optimization.get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+    )
+
+    logger.info("Preparing model, optimizer and dataloaders")
+
+    if args.train_text_encoder:
+        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+            model, optimizer, lr_scheduler, train_dataloader, text_encoder[0], text_encoder[1] = accelerator.prepare(
+                model, optimizer, lr_scheduler, train_dataloader, text_encoder[0], text_encoder[1]
+            )
+        else:
+            model, optimizer, lr_scheduler, train_dataloader, text_encoder = accelerator.prepare(
+            model, optimizer, lr_scheduler, train_dataloader, text_encoder
+        )
+    else:
+        model, optimizer, lr_scheduler, train_dataloader = accelerator.prepare(
+            model, optimizer, lr_scheduler, train_dataloader
+        )
+
+    train_dataloader.num_batches = len(train_dataloader)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    if not args.train_text_encoder:
+        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+            text_encoder[0].to(device=accelerator.device, dtype=weight_dtype)
+            text_encoder[1].to(device=accelerator.device, dtype=weight_dtype)
+        else:
+            text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+
+    vq_model.to(device=accelerator.device)
+
+    if args.use_ema:
+        ema.to(accelerator.device)
+
+    with nullcontext() if args.train_text_encoder else torch.no_grad():
+        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+            _input_ids_tmp_ = tokenize_prompt(tokenizer, "", args.text_encoder_architecture)
+            _input_ids_tmp_[0] = _input_ids_tmp_[0].to(accelerator.device, non_blocking=True)
+            _input_ids_tmp_[1] = _input_ids_tmp_[1].to(accelerator.device, non_blocking=True)
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                text_encoder, _input_ids_tmp_, args.text_encoder_architecture
+            )
+        else:
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                text_encoder, tokenize_prompt(tokenizer, "", args.text_encoder_architecture).to(accelerator.device, non_blocking=True), args.text_encoder_architecture
+            )
+
+        # There is a single image, we can just pre-encode the single prompt
+        if args.instance_data_image is not None:
+            prompt = os.path.splitext(os.path.basename(args.instance_data_image))[0]
+            if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+                _input_ids_tmp_ = tokenize_prompt(tokenizer, prompt, args.text_encoder_architecture)
+                _input_ids_tmp_[0] = _input_ids_tmp_[0].to(accelerator.device, non_blocking=True)
+                _input_ids_tmp_[1] = _input_ids_tmp_[1].to(accelerator.device, non_blocking=True)
+                empty_embeds, empty_clip_embeds = encode_prompt(
+                    text_encoder, _input_ids_tmp_, args.text_encoder_architecture
+                )
+            else:
+                encoder_hidden_states, cond_embeds = encode_prompt(
+                    text_encoder, tokenize_prompt(tokenizer, prompt, args.text_encoder_architecture).to(accelerator.device, non_blocking=True), args.text_encoder_architecture
+                )
+            encoder_hidden_states = encoder_hidden_states.repeat(args.train_batch_size, 1, 1)
+            cond_embeds = cond_embeds.repeat(args.train_batch_size, 1)
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(train_dataloader.num_batches / args.gradient_accumulation_steps)
+    # Afterwards we recalculate our number of training epochs.
+    # Note: We are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info(f"  Num training steps = {args.max_train_steps}")
+    logger.info(f"  Instantaneous batch size per device = { args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+
+    resume_from_checkpoint = args.resume_from_checkpoint
+    if resume_from_checkpoint:
+        if resume_from_checkpoint == "latest":
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            if len(dirs) > 0:
+                resume_from_checkpoint = os.path.join(args.output_dir, dirs[-1])
+            else:
+                resume_from_checkpoint = None
+
+        if resume_from_checkpoint is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+        else:
+            accelerator.print(f"Resuming from checkpoint {resume_from_checkpoint}")
+
+    if resume_from_checkpoint is None:
+        global_step = 0
+        first_epoch = 0
+    else:
+        accelerator.load_state(resume_from_checkpoint)
+        global_step = int(os.path.basename(resume_from_checkpoint).split("-")[1])
+        first_epoch = global_step // num_update_steps_per_epoch
+
+    # This is to solve the inconsistent tensor device issue
+    if args.use_ema:
+        ema.shadow_params = [p.to(accelerator.device) for p in ema.shadow_params]
+
+    # As stated above, we are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    for epoch in range(first_epoch, num_train_epochs):
+        for batch in train_dataloader:
+            torch.cuda.empty_cache()
+            with torch.no_grad():
+                micro_conds = batch["micro_conds"].to(accelerator.device, non_blocking=True)
+                pixel_values = batch["image"].to(accelerator.device, non_blocking=True)
+
+                batch_size = pixel_values.shape[0]
+
+                split_batch_size = args.split_vae_encode if args.split_vae_encode is not None else batch_size
+                num_splits = math.ceil(batch_size / split_batch_size)
+                image_tokens = []
+                for i in range(num_splits):
+                    start_idx = i * split_batch_size
+                    end_idx = min((i + 1) * split_batch_size, batch_size)
+                    bs = pixel_values.shape[0]
+                    image_tokens.append(
+                        vq_model.quantize(
+                            vq_model.encode(
+                                pixel_values[start_idx: end_idx]
+                            ).latents
+                        )[2][2].reshape(split_batch_size, -1)
+                    )
+                image_tokens = torch.cat(image_tokens, dim=0)
+
+                batch_size, seq_len = image_tokens.shape
+
+                timesteps = torch.ones(batch_size, device=image_tokens.device)
+                mask_prob = torch.cos(timesteps * math.pi * 0.5)
+                mask_prob = mask_prob.clip(args.min_masking_rate)
+
+                num_token_masked = (seq_len * mask_prob).round().clamp(min=1)
+                batch_randperm = torch.rand(batch_size, seq_len, device=image_tokens.device).argsort(dim=-1)
+                mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                mask_id = accelerator.unwrap_model(model).config.vocab_size - 1
+                input_ids = torch.where(mask, mask_id, image_tokens)
+                # labels = torch.where(mask, image_tokens, -100)
+
+                if "prompt_input_ids" in batch:
+                    prompt_input_ids = batch["prompt_input_ids"]
+                    labels = prompt_input_ids
+                    with nullcontext() if args.train_text_encoder else torch.no_grad():
+                        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+                            batch["prompt_input_ids"][0] = batch["prompt_input_ids"][0].to(accelerator.device, non_blocking=True)
+                            batch["prompt_input_ids"][1] = batch["prompt_input_ids"][1].to(accelerator.device, non_blocking=True)
+                            encoder_hidden_states, cond_embeds = encode_prompt(
+                                text_encoder, batch["prompt_input_ids"], args.text_encoder_architecture
+                            )
+                        else:
+                            encoder_hidden_states, cond_embeds = encode_prompt(
+                                text_encoder, batch["prompt_input_ids"].to(accelerator.device, non_blocking=True), args.text_encoder_architecture
+                            )
+
+                if args.cond_dropout_prob > 0.0:
+                    assert encoder_hidden_states is not None
+
+                    batch_size = encoder_hidden_states.shape[0]
+
+                    mask = (
+                        torch.zeros((batch_size, 1, 1), device=encoder_hidden_states.device).float().uniform_(0, 1)
+                        < args.cond_dropout_prob
+                    )
+
+                    empty_embeds_ = empty_embeds.expand(batch_size, -1, -1)
+                    encoder_hidden_states = torch.where(
+                        (encoder_hidden_states * mask).bool(), encoder_hidden_states, empty_embeds_
+                    )
+
+                    empty_clip_embeds_ = empty_clip_embeds.expand(batch_size, -1)
+                    cond_embeds = torch.where((cond_embeds * mask.squeeze(-1)).bool(), cond_embeds, empty_clip_embeds_)
+
+                bs = input_ids.shape[0]
+                vae_scale_factor = 2 ** (len(vq_model.config.block_out_channels) - 1)
+                resolution = args.resolution // vae_scale_factor
+                input_ids = input_ids.reshape(bs, resolution, resolution)
+
+            # Train Step
+            with accelerator.accumulate(model):
+                if args.pretrained_model_architecture == 'Meissonic':
+                   
+                    if args.resolution == 1024: # only stage 3 and stage 4 do not apply 2*
+                        img_ids = _prepare_latent_image_ids(input_ids.shape[0], input_ids.shape[-2], input_ids.shape[-1], input_ids.device, input_ids.dtype)
+                    else:
+                        img_ids = _prepare_latent_image_ids(input_ids.shape[0], 2 * input_ids.shape[-2], 2 * input_ids.shape[-1], input_ids.device, input_ids.dtype)
+
+                    txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(device = input_ids.device, dtype = input_ids.dtype)
+                   
+                    logits = model(
+                        hidden_states=input_ids, # should be (batch size, channel, height, width)
+                        encoder_hidden_states=encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                        micro_conds=micro_conds, # 
+                        pooled_projections=cond_embeds, # should be (batch_size, projection_dim)
+                        img_ids=img_ids,
+                        txt_ids=txt_ids,
+                        timestep=mask_prob * 1000,
+                    )[1]
+                    # print(logits.shape)
+                    logits = logits.reshape(-1, tokenizer.vocab_size)
+                    
+                else:
+                    raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")
+
+                loss = F.cross_entropy(
+                    logits,
+                    labels.view(-1),
+                    ignore_index=-100,
+                    reduction="mean",
+                )
+
+                # Gather the losses across all processes for logging (if we use distributed training).
+                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                avg_masking_rate = accelerator.gather(mask_prob.repeat(args.train_batch_size)).mean()
+
+                accelerator.backward(loss)
+
+                if args.max_grad_norm is not None and accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                if args.use_ema:
+                    ema.step(model.parameters())
+
+                if (global_step + 1) % args.logging_steps == 0:
+                    logs = {
+                        "step_loss": avg_loss.item(),
+                        "lr": lr_scheduler.get_last_lr()[0],
+                        "avg_masking_rate": avg_masking_rate.item(),
+                    }
+                    accelerator.log(logs, step=global_step + 1)
+
+                    logger.info(
+                        f"Step: {global_step + 1} "
+                        f"Loss: {avg_loss.item():0.4f} "
+                        f"LR: {lr_scheduler.get_last_lr()[0]:0.6f}"
+                    )
+
+                if (global_step + 1) % args.checkpointing_steps == 0:
+                    save_checkpoint(args, accelerator, global_step + 1, logger)
+
+                if (global_step + 1) % args.validation_steps == 0 and accelerator.is_main_process:
+                    if args.use_ema:
+                        ema.store(model.parameters())
+                        ema.copy_to(model.parameters())
+
+                    with torch.no_grad():
+                        logger.info("Generating images...")
+
+                        model.eval()
+
+                        scheduler = Scheduler.from_pretrained(
+                            args.pretrained_model_name_or_path,
+                            subfolder="scheduler",
+                            revision=args.revision,
+                            variant=args.variant,
+                        )
+
+                        pipe = UnifiedPipeline(
+                            transformer=accelerator.unwrap_model(model),
+                            tokenizer=tokenizer,
+                            text_encoder=text_encoder,
+                            vqvae=vq_model,
+                            scheduler=scheduler,
+                        )
+
+                        output = pipe(
+                            prompt=args.validation_prompts, 
+                            height=args.resolution, 
+                            width=args.resolution, 
+                            guidance_scale=9, 
+                            num_inference_steps=64
+                        )
+                        pil_images = output.images
+                        prompts = output.prompts
+                        print(prompts)
+
+                        result=[]
+                        for img in pil_images:
+                            if not isinstance(img, torch.Tensor):
+                                img = transforms.ToTensor()(img)
+                            result.append(img.unsqueeze(0))
+                        result = torch.cat(result,dim=0)
+                        result = make_grid(result, nrow=3)
+                        save_image(result,os.path.join(args.output_dir, str(global_step)+'_text2image_1024_CFG-9.png'))
+                        
+                        # 保存为JSON
+                        output_data = {
+                            "step": global_step,
+                            "prompts": prompts,
+                            "images": [f"{global_step}_text2image_1024_CFG-9_{i}.png" for i in range(len(pil_images))]
+                        }
+
+                        with open(os.path.join(args.output_dir, f"prompts_{global_step}.json"), "w") as f:
+                            json.dump(output_data, f, indent=2)
+                        
+                        model.train()
+
+                        if args.train_text_encoder:
+                            if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+                                text_encoder[0].train()
+                                text_encoder[1].trian()
+                            else:
+                                text_encoder.train()
+
+                    if args.use_ema:
+                        ema.restore(model.parameters())
+
+                global_step += 1
+
+            # Stop training if max steps is reached
+            if global_step >= args.max_train_steps:
+                break
+        # End for
+
+    accelerator.wait_for_everyone()
+
+    # Evaluate and save checkpoint at the end of training
+    save_checkpoint(args, accelerator, global_step, logger)
+
+    # Save the final trained checkpoint
+    if accelerator.is_main_process:
+        model = accelerator.unwrap_model(model)
+        if args.use_ema:
+            ema.copy_to(model.parameters())
+        model.save_pretrained(args.output_dir)
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main(parse_args())

train/train_text_decoder.sh

@@ -0,0 +1,31 @@
+# bash it in root path
+PYTHON_PATH='./' accelerate launch --multi_gpu --gpu_ids '2,3' --main_process_port 25001 --num_processes 2 train/train_text_decoder.py \
+        --output_dir  "./outputs/football/" \
+        --train_batch_size 16 \
+        --gradient_accumulation_steps 2 \
+        --learning_rate 1e-5 \
+        --max_grad_norm 10 \
+        --pretrained_model_name_or_path "meissonflow/meissonic" \
+        --text_encoder_architecture 'open_clip' \
+        --pretrained_model_architecture 'Meissonic' \
+        --instance_dataset 'MyParquetDataset' \
+        --instance_data_dir '/data/sqy/0000/' \
+        --resolution 1024 \
+        --mixed_precision fp16 \
+        --lr_scheduler constant \
+        --use_8bit_adam \
+        --dataloader_num_workers 0 \
+        --validation_prompts \
+            'a boy' \
+            'A serene mountain landscape with towering snow-capped peaks, a crystal-clear blue lake reflecting the mountains, dense pine forests, and a vibrant orange sunrise illuminating the sky.' \
+            'A playful golden retriever puppy with a shiny coat, bounding through a meadow filled with colorful wildflowers, under a bright, clear blue sky.' \
+            'A bustling city street at night, illuminated by vibrant neon signs in various colors, with busy pedestrians, street vendors, and a light rain creating reflective puddles on the pavement.' \
+            'A majestic, medieval castle perched on a rugged cliffside, overlooking a vast, calm ocean at sunset, with the sky painted in hues of pink, orange, and purple.' \
+            'An elegant ballerina in a white tutu, dancing gracefully on a grand stage with ornate, gold-trimmed curtains, under a spotlight that casts a soft glow.' \
+            'A cozy, rustic log cabin nestled in a snow-covered forest, with smoke rising from the stone chimney, warm lights glowing from the windows, and a path of footprints leading to the front door.'\
+            'A Cute Cat' \
+            'A Snow Mountain'\
+        --max_train_steps 30000 \
+        --checkpointing_steps 1000 \
+        --validation_steps 100 \
+        --logging_steps 10

train/train_text_encoder.py

@@ -0,0 +1,903 @@
+import argparse
+import copy
+import logging
+import math
+import os
+import json
+from contextlib import nullcontext
+from pathlib import Path
+import sys
+sys.path.append(os.getcwd())
+import gc
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from peft import LoraConfig
+from peft.utils import get_peft_model_state_dict
+
+from torch.utils.data import DataLoader, default_collate
+from torchvision import transforms
+from torchvision.utils import save_image,make_grid
+
+from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
+from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+)
+
+import diffusers.optimization
+from diffusers import EMAModel, VQModel 
+from diffusers.loaders import LoraLoaderMixin
+from diffusers.utils import is_wandb_available
+
+from src.scheduler import Scheduler
+from src.pipeline import UnifiedPipeline
+
+from train.trainer_utils import save_checkpoint
+from train.dataset_utils import ImageCaptionLargeDataset
+from train.dataset_utils import tokenize_prompt, encode_prompt
+from src.transformer import SymmetricTransformer2DModel
+
+if is_wandb_available():
+    import wandb
+    # wandb.login(key="")
+
+logger = get_logger(__name__, log_level="INFO")
+
+import torch._dynamo
+torch._dynamo.config.verbose = True
+
+# Optionally suppress errors to fall back to eager execution
+torch._dynamo.config.suppress_errors = True
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--pretrained_transformer_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--text_encoder_architecture",
+        type=str,
+        default="open_clip",
+        required=False,
+        help="The architecture of the text encoder. One of ['CLIP', 'open_clip', 'flan-t5-base','Qwen2-0.5B','gemini-2b', 'CLIP_T5_base']",
+    )
+    parser.add_argument(
+        "--text_encoder_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--instance_dataset",
+        type=str,
+        default=None,
+        required=False,
+        help="The dataset to use for training. One of ['MSCOCO600K', 'PickaPicV2']",
+    )
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--training_from_scratch",
+        type=bool,
+        default=False,
+        required=False
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--variant",
+        type=str,
+        default=None,
+        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.9999)
+    parser.add_argument("--ema_update_after_step", type=int, default=0)
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="muse_training",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
+            "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
+            "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
+            "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
+            "instructions."
+        ),
+    )
+    parser.add_argument(
+        "--logging_steps",
+        type=int,
+        default=50,
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more details"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=0.0003,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--validation_steps",
+        type=int,
+        default=100,
+        help=(
+            "Run validation every X steps. Validation consists of running the prompt"
+            " `args.validation_prompt` multiple times: `args.num_validation_images`"
+            " and logging the images."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="wandb",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument("--validation_prompts", type=str, nargs="*")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument("--split_vae_encode", type=int, required=False, default=None)
+    parser.add_argument("--min_masking_rate", type=float, default=0.0)
+    parser.add_argument("--cond_dropout_prob", type=float, default=0.0)
+    parser.add_argument("--max_grad_norm", default=50.0, type=float, help="Max gradient norm.", required=False)
+    parser.add_argument("--use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--text_encoder_use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--lora_r", default=16, type=int)
+    parser.add_argument("--lora_alpha", default=32, type=int)
+    parser.add_argument("--lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--text_encoder_lora_r", default=16, type=int)
+    parser.add_argument("--text_encoder_lora_alpha", default=32, type=int)
+    parser.add_argument("--text_encoder_lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--train_text_encoder", action="store_true")
+    parser.add_argument("--image_key", type=str, required=False)
+    parser.add_argument("--prompt_key", type=str, required=False)
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--prompt_prefix", type=str, required=False, default=None)
+
+    args = parser.parse_args()
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+    if args.instance_data_dir is not None:
+        if not os.path.exists(args.instance_data_dir):
+            raise ValueError(f"Does not exist: `--args.instance_data_dir` {args.instance_data_dir}")
+
+    return args
+
+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
+    )
+    # latent_image_ids = latent_image_ids.unsqueeze(0).repeat(batch_size, 1, 1)
+
+    return latent_image_ids.to(device=device, dtype=dtype)
+
+def main(args):
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if accelerator.is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers("muddit", config=vars(copy.deepcopy(args)))
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    if args.text_encoder_architecture == "gemma":
+        text_encoder_one = CLIPTextModelWithProjection.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
+        )
+        tokenizer_one = CLIPTokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
+        )        
+        
+        text_encoder_two = Gemma2Model.from_pretrained(
+            args.text_encoder_name_or_path, variant=args.variant
+        )
+        tokenizer_two = GemmaTokenizerFast.from_pretrained(
+            args.text_encoder_name_or_path, variant=args.variant
+        )
+        t5_dim = text_encoder_two.config.hidden_size
+        
+        text_encoder = [text_encoder_one, text_encoder_two]
+        tokenizer = [tokenizer_one, tokenizer_two]
+        
+        text_encoder_one.requires_grad_(False)
+        text_encoder_two.requires_grad_(False)
+    else:
+        raise ValueError(f"Unknown text encoder architecture: {args.text_encoder_architecture}")
+    
+    vq_model = VQModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vqvae", revision=args.revision, variant=args.variant
+    )
+    vq_model.requires_grad_(False)
+
+    model = SymmetricTransformer2DModel.from_pretrained(
+        args.pretrained_model_name_or_path if args.pretrained_transformer_path is None else args.pretrained_transformer_path, 
+        subfolder="transformer",
+        low_cpu_mem_usage=False, 
+        device_map=None,
+    )
+    
+    if args.pretrained_transformer_path is None and model.adapter is None:
+        model.register_to_config(t5_dim=t5_dim)
+        model.adapter = nn.Sequential(
+            nn.LayerNorm(t5_dim, elementwise_affine=False, eps=1e-6),
+            nn.Linear(t5_dim, model.config.joint_attention_dim, bias=False)
+        )
+   
+    model.requires_grad_(True)                
+    model.train()
+    model = torch.compile(model)
+    
+    if args.gradient_checkpointing:
+        model.enable_gradient_checkpointing()
+
+    if args.use_ema: # Not verify the robostness of this part
+        ema = EMAModel(
+            model.parameters(),
+            decay=args.ema_decay,
+            update_after_step=args.ema_update_after_step,
+            model_cls=SymmetricTransformer2DModel, 
+            model_config=model.config,
+        )
+
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            transformer_lora_layers_to_save = None
+            text_encoder_lora_layers_to_save = None
+
+            for model_ in models:
+                if isinstance(model_, type(accelerator.unwrap_model(model))):
+                    if args.use_lora:
+                        transformer_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "transformer"))
+                elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                    if args.text_encoder_use_lora:
+                        text_encoder_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "text_encoder"))
+                else:
+                    raise ValueError(f"unexpected save model: {model_.__class__}")
+
+                # make sure to pop weight so that corresponding model is not saved again
+                weights.pop()
+
+            if transformer_lora_layers_to_save is not None or text_encoder_lora_layers_to_save is not None:
+                LoraLoaderMixin.save_lora_weights(
+                    output_dir,
+                    unet_lora_layers=transformer_lora_layers_to_save,
+                    text_encoder_lora_layers=text_encoder_lora_layers_to_save,
+                )
+
+            if args.use_ema:
+                ema.save_pretrained(os.path.join(output_dir, "ema_model"))
+
+    def load_model_hook(models, input_dir):
+        transformer = None
+        text_encoder_ = None
+
+        # this part is added for keep consistency when add model.compile() in the model
+        def adap_compile(ori_dict):#add '_orig_mod.' to each key
+            new_dict = {}
+            for k,v in ori_dict.items():
+                new_dict['_orig_mod.'+k] = v
+            return new_dict
+            
+        while len(models) > 0:
+            model_ = models.pop()
+
+            if isinstance(model_, type(accelerator.unwrap_model(model))):
+                if args.use_lora:
+                    transformer = model_
+                else:
+                    load_model = SymmetricTransformer2DModel.from_pretrained(os.path.join(input_dir, "transformer"), low_cpu_mem_usage=False, device_map=None)
+                    model_.load_state_dict(adap_compile(load_model.state_dict()))
+                    del load_model
+            else:
+                raise ValueError(f"unexpected save model: {model.__class__}")
+
+        if transformer is not None or text_encoder_ is not None:
+            lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir)
+            LoraLoaderMixin.load_lora_into_text_encoder(
+                lora_state_dict, network_alphas=network_alphas, text_encoder=text_encoder_
+            )
+            LoraLoaderMixin.load_lora_into_transformer(
+                lora_state_dict, network_alphas=network_alphas, transformer=transformer
+            )
+
+        if args.use_ema:
+            load_from = EMAModel.from_pretrained(os.path.join(input_dir, "ema_model"), model_cls=SymmetricTransformer2DModel)
+            ema.load_state_dict(adap_compile(load_from.state_dict()))
+            del load_from
+
+    accelerator.register_load_state_pre_hook(load_model_hook)
+    accelerator.register_save_state_pre_hook(save_model_hook)
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        )
+
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+            )
+
+        optimizer_cls = bnb.optim.AdamW8bit
+    else:
+        optimizer_cls = torch.optim.AdamW
+
+    optimizer_parameters = [p for p in model.parameters() if p.requires_grad]
+    optimizer = optimizer_cls(
+        optimizer_parameters,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    logger.info("Creating dataloaders and lr_scheduler")
+
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    if args.instance_dataset == "ImageCaptionLargeDataset":
+        dataset = ImageCaptionLargeDataset(
+            root_dir=args.instance_data_dir,
+            tokenizer=tokenizer,
+            size=args.resolution,
+            text_encoder_architecture=args.text_encoder_architecture
+        )
+    elif args.instance_dataset == "DATA_TYPE":
+        raise NotImplementedError("DATA_TYPE is not yet supported")
+    else:
+        assert False
+
+    def collate_fn(samples):
+        images = [sample["image"] for sample in samples]
+        micro_conds = [sample["micro_conds"] for sample in samples]
+        
+        images = torch.stack(images, dim=0)
+        micro_conds = torch.stack(micro_conds, dim=0)
+        
+        if isinstance(samples[0]["prompt_input_ids"], list):
+            input_ids = [sample["prompt_input_ids"][0] for sample in samples]
+            input_ids_2 = [sample["prompt_input_ids"][1] for sample in samples]
+            
+            input_ids = torch.cat(input_ids, dim=0)
+            input_ids_2 = torch.cat(input_ids_2, dim=0)
+            prompt_input_ids = [input_ids, input_ids_2]
+            
+        elif isinstance(samples[0]["prompt_input_ids"], torch.Tensor):
+            input_ids = [sample["prompt_input_ids"] for sample in samples]
+            
+            input_ids = torch.cat(input_ids, dim=0)
+            prompt_input_ids = input_ids
+    
+        ret = dict(
+            images=images,
+            micro_conds=micro_conds,
+            prompt_input_ids=prompt_input_ids,
+        )
+        
+        return ret
+
+    train_dataloader = DataLoader(
+        dataset,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=collate_fn,
+        pin_memory=True,
+    )
+    train_dataloader.num_batches = len(train_dataloader)
+
+    lr_scheduler = diffusers.optimization.get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+    )
+
+    logger.info("Preparing model, optimizer and dataloaders")
+
+    if args.train_text_encoder:
+        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
+            model, optimizer, lr_scheduler, train_dataloader, text_encoder[0], text_encoder[1] = accelerator.prepare(
+                model, optimizer, lr_scheduler, train_dataloader, text_encoder[0], text_encoder[1]
+            )
+        else:
+            model, optimizer, lr_scheduler, train_dataloader, text_encoder = accelerator.prepare(
+                model, optimizer, lr_scheduler, train_dataloader, text_encoder
+            )
+    else:
+        model, optimizer, lr_scheduler, train_dataloader = accelerator.prepare(
+            model, optimizer, lr_scheduler, train_dataloader
+        )
+
+    train_dataloader.num_batches = len(train_dataloader)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    if not args.train_text_encoder:
+        if args.text_encoder_architecture in ("t5_clip", "gemma"): # Not support yet. Only support open_clip
+            text_encoder[0].to(device=accelerator.device, dtype=weight_dtype)
+            text_encoder[1].to(device=accelerator.device, dtype=weight_dtype)
+        else:
+            text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+
+    vq_model.to(device=accelerator.device)
+
+    if args.use_ema:
+        ema.to(accelerator.device)
+
+    with nullcontext() if args.train_text_encoder else torch.no_grad():
+        if args.text_encoder_architecture in ("t5_clip", "gemma"): # Not support yet. Only support open_clip
+            _input_ids_tmp_ = tokenize_prompt(tokenizer, "", args.text_encoder_architecture)
+            _input_ids_tmp_[0] = _input_ids_tmp_[0].to(accelerator.device, non_blocking=True)
+            _input_ids_tmp_[1] = _input_ids_tmp_[1].to(accelerator.device, non_blocking=True)
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                text_encoder, 
+                _input_ids_tmp_, 
+                args.text_encoder_architecture
+            )
+        else:
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                text_encoder, 
+                tokenize_prompt(tokenizer, "", args.text_encoder_architecture).to(accelerator.device, non_blocking=True), 
+                args.text_encoder_architecture
+            )
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(train_dataloader.num_batches / args.gradient_accumulation_steps)
+    # Afterwards we recalculate our number of training epochs.
+    # Note: We are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info(f"  Num training steps = {args.max_train_steps}")
+    logger.info(f"  Instantaneous batch size per device = { args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+
+    resume_from_checkpoint = args.resume_from_checkpoint
+    if resume_from_checkpoint:
+        if resume_from_checkpoint == "latest":
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            if len(dirs) > 0:
+                resume_from_checkpoint = os.path.join(args.output_dir, dirs[-1])
+            else:
+                resume_from_checkpoint = None
+
+        if resume_from_checkpoint is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+        else:
+            accelerator.print(f"Resuming from checkpoint {resume_from_checkpoint}")
+
+    if resume_from_checkpoint is None:
+        global_step = 0
+        first_epoch = 0
+    else:
+        accelerator.load_state(resume_from_checkpoint)
+        global_step = int(os.path.basename(resume_from_checkpoint).split("-")[1])
+        first_epoch = global_step // num_update_steps_per_epoch
+
+    # This is to solve the inconsistent tensor device issue
+    if args.use_ema:
+        ema.shadow_params = [p.to(accelerator.device) for p in ema.shadow_params]
+
+    # As stated above, we are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    for epoch in range(first_epoch, num_train_epochs):
+        for batch in train_dataloader:
+            torch.cuda.empty_cache()
+            with torch.no_grad():
+                micro_conds = batch["micro_conds"].to(accelerator.device, non_blocking=True)
+                pixel_values = batch["images"].to(accelerator.device, non_blocking=True)
+
+                batch_size = pixel_values.shape[0]
+
+                # ====================== tokenize images ======================
+                split_batch_size = args.split_vae_encode if args.split_vae_encode is not None else batch_size
+                num_splits = math.ceil(batch_size / split_batch_size)
+                image_tokens = []
+                for i in range(num_splits):
+                    start_idx = i * split_batch_size
+                    end_idx = min((i + 1) * split_batch_size, batch_size)
+                    bs = pixel_values.shape[0]
+                    image_tokens.append(
+                        vq_model.quantize(
+                            vq_model.encode(
+                                pixel_values[start_idx: end_idx]
+                            ).latents
+                        )[2][2].reshape(split_batch_size, -1)
+                    )
+                image_tokens = torch.cat(image_tokens, dim=0)
+                # ====================== tokenize images ======================
+
+                batch_size, seq_len = image_tokens.shape
+                timesteps = torch.rand(batch_size, device=image_tokens.device)
+                mask_prob = torch.cos(timesteps * math.pi * 0.5)
+                mask_prob = mask_prob.clip(args.min_masking_rate)
+
+                num_token_masked = (seq_len * mask_prob).round().clamp(min=1)
+                batch_randperm = torch.rand(batch_size, seq_len, device=image_tokens.device).argsort(dim=-1)
+                mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                mask_id = accelerator.unwrap_model(model).config.vocab_size - 1
+                input_ids = torch.where(mask, mask_id, image_tokens)
+                labels = torch.where(mask, image_tokens, -100)
+
+                if "prompt_input_ids" in batch:
+                    with nullcontext() if args.train_text_encoder else torch.no_grad():
+                        if args.text_encoder_architecture in ("t5_clip", "gemma"): # Not support yet. Only support open_clip
+                            batch["prompt_input_ids"][0] = batch["prompt_input_ids"][0].to(accelerator.device, non_blocking=True)
+                            batch["prompt_input_ids"][1] = batch["prompt_input_ids"][1].to(accelerator.device, non_blocking=True)
+                            encoder_hidden_states, cond_embeds = encode_prompt(
+                                text_encoder, 
+                                batch["prompt_input_ids"], 
+                                args.text_encoder_architecture
+                            )
+                        else:
+                            encoder_hidden_states, cond_embeds = encode_prompt(
+                                text_encoder, 
+                                batch["prompt_input_ids"].to(accelerator.device, non_blocking=True), 
+                                args.text_encoder_architecture
+                            )
+
+                if args.cond_dropout_prob > 0.0:
+                    assert encoder_hidden_states is not None
+
+                    batch_size = encoder_hidden_states.shape[0]
+
+                    mask = (
+                        torch.zeros((batch_size, 1, 1), device=encoder_hidden_states.device).float().uniform_(0, 1)
+                        < args.cond_dropout_prob
+                    )
+
+                    empty_embeds_ = empty_embeds.expand(batch_size, -1, -1)
+                    encoder_hidden_states = torch.where(
+                        (encoder_hidden_states * mask).bool(), encoder_hidden_states, empty_embeds_
+                    )
+
+                    empty_clip_embeds_ = empty_clip_embeds.expand(batch_size, -1)
+                    cond_embeds = torch.where((cond_embeds * mask.squeeze(-1)).bool(), cond_embeds, empty_clip_embeds_)
+
+                bs = input_ids.shape[0]
+                vae_scale_factor = 2 ** (len(vq_model.config.block_out_channels) - 1)
+                resolution = args.resolution // vae_scale_factor
+                input_ids = input_ids.reshape(bs, resolution, resolution)
+
+            # Train Step
+            with accelerator.accumulate(model):
+                codebook_size = accelerator.unwrap_model(model).config.codebook_size
+                if args.resolution == 1024: # only stage 3 and stage 4 do not apply 2*
+                    img_ids = _prepare_latent_image_ids(input_ids.shape[0], input_ids.shape[-2], input_ids.shape[-1], input_ids.device, input_ids.dtype)
+                else:
+                    img_ids = _prepare_latent_image_ids(input_ids.shape[0], 2 * input_ids.shape[-2], 2 * input_ids.shape[-1], input_ids.device, input_ids.dtype)
+
+                txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(device = input_ids.device, dtype = input_ids.dtype)
+                
+                logits = model(
+                    hidden_states=input_ids, # should be (batch size, channel, height, width)
+                    encoder_hidden_states=encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                    micro_conds=micro_conds, # 
+                    pooled_projections=cond_embeds, # should be (batch_size, projection_dim)
+                    img_ids=img_ids,
+                    txt_ids=txt_ids,
+                    timestep=mask_prob,
+                )[0]
+                logits = logits.reshape(batch_size, codebook_size, -1).permute(0, 2, 1)
+                logits = logits.reshape(-1, codebook_size)
+                loss = F.cross_entropy(
+                    logits,
+                    labels.view(-1),
+                    ignore_index=-100,
+                    reduction="mean",
+                )
+
+                # Gather the losses across all processes for logging (if we use distributed training).
+                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                avg_masking_rate = accelerator.gather(mask_prob.repeat(args.train_batch_size)).mean()
+
+                accelerator.backward(loss)
+
+                if args.max_grad_norm is not None and accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                if args.use_ema:
+                    ema.step(model.parameters())
+
+                if (global_step + 1) % args.logging_steps == 0:
+                    logs = {
+                        "step_loss": avg_loss.item(),
+                        "lr": lr_scheduler.get_last_lr()[0],
+                        "avg_masking_rate": avg_masking_rate.item(),
+                    }
+                    accelerator.log(logs, step=global_step + 1)
+
+                    logger.info(
+                        f"Step: {global_step + 1} "
+                        f"Loss: {avg_loss.item():0.4f} "
+                        f"LR: {lr_scheduler.get_last_lr()[0]:0.6f}"
+                    )
+
+                if (global_step + 1) % args.checkpointing_steps == 0:
+                    save_checkpoint(args, accelerator, global_step + 1, logger)
+
+                if (global_step + 1) % args.validation_steps == 0 and accelerator.is_main_process:
+                    if args.use_ema:
+                        ema.store(model.parameters())
+                        ema.copy_to(model.parameters())
+
+                    with torch.no_grad():
+                        logger.info("Generating images...")
+
+                        model.eval()
+
+                        scheduler = Scheduler.from_pretrained(
+                            args.pretrained_model_name_or_path,
+                            subfolder="scheduler",
+                            revision=args.revision,
+                            variant=args.variant,
+                        )
+
+                        pipe = UnifiedPipeline(
+                            transformer=accelerator.unwrap_model(model),
+                            tokenizer=tokenizer_one,
+                            tokenizer_2=tokenizer_two,
+                            text_encoder=text_encoder_one,
+                            text_encoder_2=text_encoder_two,
+                            vqvae=vq_model,
+                            scheduler=scheduler,
+                        )
+
+                        output = pipe(
+                            prompt=args.validation_prompts, 
+                            height=args.resolution, 
+                            width=args.resolution, 
+                            guidance_scale=9, 
+                            num_inference_steps=64
+                        )
+                        pil_images = output.images
+                        
+                        wandb_images = [
+                            wandb.Image(image, caption=args.validation_prompts[i])
+                            for i, image in enumerate(pil_images)
+                        ]
+
+                        wandb.log({"generated_images": wandb_images}, step=global_step + 1)
+                        
+                        result=[]
+                        for img in pil_images:
+                            if not isinstance(img, torch.Tensor):
+                                img = transforms.ToTensor()(img)
+                            result.append(img.unsqueeze(0))
+                        result = torch.cat(result,dim=0)
+                        result = make_grid(result, nrow=3)
+                        save_image(result,os.path.join(args.output_dir, str(global_step)+'_text2image_1024_CFG-9.png'))
+                        
+                        model.train()
+
+                    if args.use_ema:
+                        ema.restore(model.parameters())
+
+                global_step += 1
+
+            # Stop training if max steps is reached
+            if global_step >= args.max_train_steps:
+                break
+        # End for
+
+    accelerator.wait_for_everyone()
+
+    # Evaluate and save checkpoint at the end of training
+    save_checkpoint(args, accelerator, global_step, logger)
+
+    # Save the final trained checkpoint
+    if accelerator.is_main_process:
+        model = accelerator.unwrap_model(model)
+        if args.use_ema:
+            ema.copy_to(model.parameters())
+        model.save_pretrained(args.output_dir)
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main(parse_args())

train/train_text_encoder.sh

@@ -0,0 +1,34 @@
+export TOKENIZERS_PARALLELISM=false
+
+# bash it in root path
+PYTHON_PATH='./' accelerate launch --multi_gpu --gpu_ids '2,3' --main_process_port 25000 --num_processes 2 train/train_text_encoder.py \
+        --output_dir "./outputs/debug/" \
+        --train_batch_size 8 \
+        --gradient_accumulation_steps 2 \
+        --learning_rate 1e-4 \
+        --max_grad_norm 10 \
+        --pretrained_model_name_or_path "MeissonFlow/Meissonic" \
+        --text_encoder_architecture 'gemma' \
+        --text_encoder_name_or_path "google/gemma-2-2b-it" \
+        --instance_dataset 'ImageCaptionLargeDataset' \
+        --instance_data_dir '/data/sqy/0000/' \
+        --resolution 1024 \
+        --mixed_precision fp16 \
+        --lr_scheduler constant \
+        --use_8bit_adam \
+        --dataloader_num_workers 4 \
+        --validation_prompts \
+            'a boy' \
+            'A serene mountain landscape with towering snow-capped peaks, a crystal-clear blue lake reflecting the mountains, dense pine forests, and a vibrant orange sunrise illuminating the sky.' \
+            'A playful golden retriever puppy with a shiny coat, bounding through a meadow filled with colorful wildflowers, under a bright, clear blue sky.' \
+            'A bustling city street at night, illuminated by vibrant neon signs in various colors, with busy pedestrians, street vendors, and a light rain creating reflective puddles on the pavement.' \
+            'A majestic, medieval castle perched on a rugged cliffside, overlooking a vast, calm ocean at sunset, with the sky painted in hues of pink, orange, and purple.' \
+            'An elegant ballerina in a white tutu, dancing gracefully on a grand stage with ornate, gold-trimmed curtains, under a spotlight that casts a soft glow.' \
+            'A cozy, rustic log cabin nestled in a snow-covered forest, with smoke rising from the stone chimney, warm lights glowing from the windows, and a path of footprints leading to the front door.'\
+            'A Cute Cat' \
+            'A Snow Mountain'\
+        --max_train_steps 30000 \
+        --checkpointing_steps 1000 \
+        --validation_steps 100 \
+        --logging_steps 10 \
+        --report_to "wandb"

train/train_unified.py

@@ -0,0 +1,1141 @@
+# 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 argparse
+import copy
+import logging
+import math
+import os
+from pathlib import Path
+import sys
+sys.path.append(os.getcwd())
+import json
+import gc
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from peft import LoraConfig
+from peft.utils import get_peft_model_state_dict
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
+from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    T5EncoderModel,
+    T5Tokenizer,
+)
+
+import diffusers.optimization
+from diffusers import VQModel
+
+from src.scheduler import Scheduler
+from diffusers.loaders import LoraLoaderMixin
+from diffusers.utils import is_wandb_available
+from src.pipeline import UnifiedPipeline
+from torchvision.utils import save_image, make_grid
+from train.trainer_utils import save_checkpoint
+from train.dataset_utils import ImageCaptionLargeDataset
+from train.dataset_utils import tokenize_prompt, encode_prompt
+from src.transformer import Transformer2DModel, SymmetricTransformer2DModel
+from train.trainer_utils import load_images_to_tensor
+
+if is_wandb_available():
+    import wandb
+    # wandb.login(key="")
+
+logger = get_logger(__name__, log_level="INFO")
+
+import torch._dynamo
+torch._dynamo.config.verbose = True
+
+# Optionally suppress errors to fall back to eager execution
+torch._dynamo.config.suppress_errors = True
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--pretrained_transformer_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--text_encoder_architecture",
+        type=str,
+        default="open_clip",
+        required=False,
+        help="The architecture of the text encoder. One of ['CLIP', 'open_clip', 'flan-t5-base','Qwen2-0.5B','gemini-2b',long_t5_clip','t5_clip']",
+    )
+    parser.add_argument(
+        "--text_encoder_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--remove_pooled_embeddings",
+        type=bool,
+        default=False,
+        required=False,
+        help="Whether to remove the pooled embeddings from the text encoder.",
+    )
+    parser.add_argument(
+        "--instance_dataset",
+        type=str,
+        default=None,
+        required=False,
+        help="The dataset to use for training. One of ['MSCOCO600K', 'PickaPicV2']",
+    )
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--training_from_scratch",
+        type=bool,
+        default=False,
+        required=False
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--variant",
+        type=str,
+        default=None,
+        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+    )
+    parser.add_argument(
+        "--instance_data_image", type=str, default=None, required=False, help="A single training image"
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.9999)
+    parser.add_argument("--ema_update_after_step", type=int, default=0)
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="muse_training",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
+            "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
+            "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
+            "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
+            "instructions."
+        ),
+    )
+    parser.add_argument(
+        "--logging_steps",
+        type=int,
+        default=50,
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more details"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--text_loss_weight",
+        type=float,
+        default=0.2,
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=0.0003,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--validation_steps",
+        type=int,
+        default=100,
+        help=(
+            "Run validation every X steps. Validation consists of running the prompt"
+            " `args.validation_prompt` multiple times: `args.num_validation_images`"
+            " and logging the images."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="wandb",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument("--validation_prompts", type=str, nargs="*")
+    parser.add_argument("--validation_images", type=str, default="./validation_images")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument("--split_vae_encode", type=int, required=False, default=None)
+    parser.add_argument("--min_masking_rate", type=float, default=0.0)
+    parser.add_argument("--cond_dropout_prob", type=float, default=0.0)
+    parser.add_argument("--max_grad_norm", default=50.0, type=float, help="Max gradient norm.", required=False)
+    parser.add_argument("--use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--text_encoder_use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--lora_r", default=16, type=int)
+    parser.add_argument("--lora_alpha", default=32, type=int)
+    parser.add_argument("--lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--text_encoder_lora_r", default=16, type=int)
+    parser.add_argument("--text_encoder_lora_alpha", default=32, type=int)
+    parser.add_argument("--text_encoder_lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--train_text_encoder", action="store_true")
+    parser.add_argument("--image_to_text_only", action="store_true")
+    parser.add_argument("--image_key", type=str, required=False)
+    parser.add_argument("--prompt_key", type=str, required=False)
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--prompt_prefix", type=str, required=False, default=None)
+
+    args = parser.parse_args()
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+    num_datasources = sum(
+        [x is not None for x in [args.instance_data_dir, args.instance_data_image]]
+    )
+
+    if num_datasources != 1:
+        raise ValueError(
+            "provide one and only one of `--instance_data_dir`, `--instance_data_image`, or `--instance_data_dataset`"
+        )
+
+    if args.instance_data_dir is not None:
+        if not os.path.exists(args.instance_data_dir):
+            raise ValueError(f"Does not exist: `--args.instance_data_dir` {args.instance_data_dir}")
+
+    if args.instance_data_image is not None:
+        if not os.path.exists(args.instance_data_image):
+            raise ValueError(f"Does not exist: `--args.instance_data_image` {args.instance_data_image}")
+
+    return args
+
+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 main(args):
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if accelerator.is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+
+    if accelerator.is_main_process:
+        accelerator.init_trackers("meissonic", config=vars(copy.deepcopy(args)))
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    if args.text_encoder_architecture == "open_clip":
+        text_encoder = CLIPTextModelWithProjection.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
+        )
+        tokenizer = CLIPTokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
+        )
+        tokenizer_2 = None
+        text_encoder_2 = None
+
+        extra_id_0_token = "<extra_id_0>"
+        num_new_tokens = tokenizer.add_tokens(extra_id_0_token)
+        mask_id_1 = tokenizer.convert_tokens_to_ids(extra_id_0_token)
+        if num_new_tokens > 0:
+            text_encoder.resize_token_embeddings(len(tokenizer))
+            mask_token_embedding = text_encoder.get_input_embeddings().weight[mask_id_1]
+            mask_token_embedding = mask_token_embedding.clone().detach().cpu().float()
+            if accelerator.is_main_process:
+                print("Saving masked token embedding...")
+                torch.save(mask_token_embedding, os.path.join(args.output_dir, "mask_token_embedding.pth"))
+                            
+        text_encoder.requires_grad_(False)
+    elif args.text_encoder_architecture == "t5_clip":
+        tokenizer = CLIPTokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
+        )
+        text_encoder = CLIPTextModelWithProjection.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
+        )   
+        
+        tokenizer_2 = T5Tokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer_2", variant=args.variant,
+        )
+        text_encoder_2 = T5EncoderModel.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder_2", variant=args.variant,
+        )
+        
+        text_encoder.requires_grad_(False)
+        text_encoder_2.requires_grad_(False)
+    elif args.text_encoder_architecture == "gemma":
+        tokenizer = CLIPTokenizer.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
+        )
+        text_encoder = CLIPTextModelWithProjection.from_pretrained(
+            args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
+        )   
+        
+        tokenizer_2 = GemmaTokenizerFast.from_pretrained(
+            args.text_encoder_name_or_path, variant=args.variant,
+        )
+        text_encoder_2 = Gemma2Model.from_pretrained(
+            args.text_encoder_name_or_path, variant=args.variant,
+        )
+
+        extra_id_0_token = "<extra_id_0>"
+        
+        tokenizer.add_tokens(extra_id_0_token)
+        tokenizer_2.add_tokens(extra_id_0_token)
+        mask_id_1 = tokenizer.convert_tokens_to_ids(extra_id_0_token)
+        mask_id_2 = tokenizer_2.convert_tokens_to_ids(extra_id_0_token)
+        
+        text_encoder.resize_token_embeddings(len(tokenizer))
+        text_encoder_2.resize_token_embeddings(len(tokenizer_2))
+        
+        text_encoder.requires_grad_(False)
+        text_encoder_2.requires_grad_(False)
+    else:
+        raise ValueError(f"Unknown text encoder architecture: {args.text_encoder_architecture}")
+    
+    vq_model = VQModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vqvae", revision=args.revision, variant=args.variant
+    )
+    vq_model.requires_grad_(False)
+
+    model = SymmetricTransformer2DModel.from_pretrained(
+        args.pretrained_transformer_path, 
+        subfolder="transformer", 
+        low_cpu_mem_usage=False, 
+        device_map=None
+    )
+
+    if model.config.tokenizer_vocab_size is None:
+        if args.text_encoder_architecture == "open_clip":
+            model.register_to_config(tokenizer_vocab_size=len(tokenizer))
+        elif args.text_encoder_architecture in ("t5_clip", "gemma"):
+            model.register_to_config(tokenizer_vocab_size=len(tokenizer_2))
+            if model.adapter is None:
+                raise ValueError(f"The MMDiT must has adapter if you want to use t5_clip mode!!!")
+        else:
+            raise ValueError(f"Unknown text encoder architecture!")
+
+        if accelerator.is_main_process:
+            print(f"model's tokenizer vocab size is {model.config.tokenizer_vocab_size}")
+            
+        model.text_decoder = nn.Sequential(
+            nn.LayerNorm(model.inner_dim, elementwise_affine=False, eps=1e-6),
+            nn.Linear(model.inner_dim, model.config.tokenizer_vocab_size, bias=False)
+        )
+
+    model = torch.compile(model)
+
+    if args.use_lora:
+        lora_config = LoraConfig(
+            r=args.lora_r,
+            lora_alpha=args.lora_alpha,
+            target_modules=args.lora_target_modules,
+        )
+        model.add_adapter(lora_config)
+
+    model.train()
+
+    if args.image_to_text_only:
+        frozen_keys = ["project_from_hidden", "up_block", "mlm_layer"]
+        for n, p in model.named_parameters():
+            if any([frozen_key in n for frozen_key in frozen_keys]):
+                p.requires_grad_(False)
+            else:
+                p.requires_grad_(True)
+    else:
+        model.requires_grad_(True)
+
+    if args.gradient_checkpointing:
+        model.enable_gradient_checkpointing()   
+
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            transformer_lora_layers_to_save = None
+            text_encoder_lora_layers_to_save = None
+
+            for model_ in models:
+                if isinstance(model_, type(accelerator.unwrap_model(model))):
+                    if args.use_lora:
+                        transformer_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "transformer"))
+                elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                    if args.text_encoder_use_lora:
+                        text_encoder_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "text_encoder"))
+                else:
+                    raise ValueError(f"unexpected save model: {model_.__class__}")
+
+                # make sure to pop weight so that corresponding model is not saved again
+                weights.pop()
+
+            if transformer_lora_layers_to_save is not None or text_encoder_lora_layers_to_save is not None:
+                LoraLoaderMixin.save_lora_weights(
+                    output_dir,
+                    unet_lora_layers=transformer_lora_layers_to_save,
+                    text_encoder_lora_layers=text_encoder_lora_layers_to_save,
+                )
+                
+
+    def load_model_hook(models, input_dir):
+        transformer = None
+        text_encoder_ = None
+
+        # this part is added for keep consistency when add model.compile() in the model
+        def adap_compile(ori_dict):#add '_orig_mod.' to each key
+            new_dict = {}
+            for k,v in ori_dict.items():
+                new_dict['_orig_mod.' + k] = v
+            return new_dict
+            
+        while len(models) > 0:
+            model_ = models.pop()
+
+            if isinstance(model_, type(accelerator.unwrap_model(model))):
+                if args.use_lora:
+                    transformer = model_
+                else:
+                    load_model = SymmetricTransformer2DModel.from_pretrained(os.path.join(input_dir, "transformer"), low_cpu_mem_usage=False, device_map=None)
+                    model_.load_state_dict(adap_compile(load_model.state_dict()))
+                    del load_model
+            elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                if args.text_encoder_use_lora:
+                    text_encoder_ = model_
+                else:
+                    try:
+                        load_model = CLIPTextModelWithProjection.from_pretrained(os.path.join(input_dir, "text_encoder"))
+                        model_.load_state_dict(load_model.state_dict())
+                        # print('finished loading text encoder!')
+                    except:
+                        print('Not found text-encoder model in current folder. So we download one text encoder from Internet.')
+                        load_model = CLIPTextModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
+                        model_.load_state_dict(load_model.state_dict())
+                    del load_model
+            else:
+                raise ValueError(f"unexpected save model: {model.__class__}")
+
+        if transformer is not None or text_encoder_ is not None:
+            lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir)
+            LoraLoaderMixin.load_lora_into_text_encoder(
+                lora_state_dict, network_alphas=network_alphas, text_encoder=text_encoder_
+            )
+            LoraLoaderMixin.load_lora_into_transformer(
+                lora_state_dict, network_alphas=network_alphas, transformer=transformer
+            )
+
+    accelerator.register_load_state_pre_hook(load_model_hook)
+    accelerator.register_save_state_pre_hook(save_model_hook)
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        )
+
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+            )
+
+        optimizer_cls = bnb.optim.AdamW8bit
+    else:
+        optimizer_cls = torch.optim.AdamW
+
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for p in model.parameters() if p.requires_grad],
+            "weight_decay": args.adam_weight_decay,
+        }
+    ]
+    optimizer = optimizer_cls(
+        optimizer_grouped_parameters,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    logger.info("Creating dataloaders and lr_scheduler")
+
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    if args.text_encoder_architecture in ("t5_clip", "gemma"):
+        tokenizer_for_dataset = [tokenizer, tokenizer_2]
+    else:
+        tokenizer_for_dataset = tokenizer
+
+    if args.instance_dataset == "ImageCaptionLargeDataset":
+        dataset = ImageCaptionLargeDataset(
+            root_dir=args.instance_data_dir,
+            tokenizer=tokenizer_for_dataset,
+            size=args.resolution,
+            text_encoder_architecture=args.text_encoder_architecture
+        )
+    elif args.instance_dataset == "DATA_TYPE":
+        raise NotImplementedError("DATA_TYPE is not yet supported")
+    else:
+        assert False
+
+    def collate_fn(samples):
+        images = [sample["image"] for sample in samples]
+        micro_conds = [sample["micro_conds"] for sample in samples]
+        
+        images = torch.stack(images, dim=0)
+        micro_conds = torch.stack(micro_conds, dim=0)
+        
+        if isinstance(samples[0]["prompt_input_ids"], list):
+            input_ids = [sample["prompt_input_ids"][0] for sample in samples]
+            input_ids_2 = [sample["prompt_input_ids"][1] for sample in samples]
+            
+            input_ids = torch.cat(input_ids, dim=0)
+            input_ids_2 = torch.cat(input_ids_2, dim=0)
+            prompt_input_ids = [input_ids, input_ids_2]
+        else:
+            input_ids = [sample["prompt_input_ids"] for sample in samples]
+            
+            input_ids = torch.cat(input_ids, dim=0)
+            prompt_input_ids = input_ids
+        
+        ret = dict(
+            images=images,
+            micro_conds=micro_conds,
+            prompt_input_ids=prompt_input_ids,
+        )
+        
+        return ret
+    
+    train_dataloader = DataLoader(
+        dataset,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=collate_fn,
+        pin_memory=True,
+    )
+    train_dataloader.num_batches = len(train_dataloader)
+
+    lr_scheduler = diffusers.optimization.get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+    )
+
+    logger.info("Preparing model, optimizer and dataloaders")
+
+    model, optimizer, lr_scheduler, train_dataloader = accelerator.prepare(
+        model, optimizer, lr_scheduler, train_dataloader
+    )
+
+    train_dataloader.num_batches = len(train_dataloader)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    if args.text_encoder_architecture in ("t5_clip", "gemma"):
+        text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+        text_encoder_2.to(device=accelerator.device, dtype=weight_dtype)
+    else:
+        text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+            
+    vq_model.to(device=accelerator.device)
+
+    with torch.no_grad():
+        if args.text_encoder_architecture in ("t5_clip", "gemma"):
+            _input_ids_tmp_ = tokenize_prompt([tokenizer, tokenizer_2], "", args.text_encoder_architecture)
+            _input_ids_tmp_[0] = _input_ids_tmp_[0].to(accelerator.device)
+            _input_ids_tmp_[1] = _input_ids_tmp_[1].to(accelerator.device)
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                [text_encoder, text_encoder_2], 
+                _input_ids_tmp_, 
+                args.text_encoder_architecture
+            )
+        else:
+            _input_ids_tmp_ = tokenize_prompt(tokenizer, "", args.text_encoder_architecture)
+            _input_ids_tmp_ = _input_ids_tmp_.to(accelerator.device)
+            empty_embeds, empty_clip_embeds = encode_prompt(
+                text_encoder,
+                _input_ids_tmp_,
+                args.text_encoder_architecture
+            )
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(train_dataloader.num_batches / args.gradient_accumulation_steps)
+    # Afterwards we recalculate our number of training epochs.
+    # Note: We are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info(f"  Num training steps = {args.max_train_steps}")
+    logger.info(f"  Instantaneous batch size per device = { args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+
+    resume_from_checkpoint = args.resume_from_checkpoint
+    if resume_from_checkpoint:
+        if resume_from_checkpoint == "latest":
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            if len(dirs) > 0:
+                resume_from_checkpoint = os.path.join(args.output_dir, dirs[-1])
+            else:
+                resume_from_checkpoint = None
+
+        if resume_from_checkpoint is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+        else:
+            accelerator.print(f"Resuming from checkpoint {resume_from_checkpoint}")
+
+    if resume_from_checkpoint is None:
+        global_step = 0
+        first_epoch = 0
+    else:
+        accelerator.load_state(resume_from_checkpoint)
+        global_step = int(os.path.basename(resume_from_checkpoint).split("-")[1])
+        first_epoch = global_step // num_update_steps_per_epoch
+
+    # As stated above, we are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    for epoch in range(first_epoch, num_train_epochs):
+        for batch in train_dataloader:
+            torch.cuda.empty_cache()
+            with torch.no_grad():
+                micro_conds = batch["micro_conds"].to(accelerator.device, non_blocking=True)
+                image_micro_conds, text_micro_conds = micro_conds.chunk(2, dim=0)
+                
+                pixel_values = batch["images"].to(accelerator.device, non_blocking=True)
+                batch_size = pixel_values.shape[0]
+
+                # ====================== tokenize images ======================
+                split_batch_size = args.split_vae_encode if args.split_vae_encode is not None else batch_size
+                num_splits = math.ceil(batch_size / split_batch_size)
+                image_tokens = []
+                for i in range(num_splits):
+                    start_idx = i * split_batch_size
+                    end_idx = min((i + 1) * split_batch_size, batch_size)
+                    image_tokens.append(
+                        vq_model.quantize(
+                            vq_model.encode(pixel_values[start_idx:end_idx]).latents
+                        )[2][2].reshape(split_batch_size, -1)
+                    )
+                image_tokens = torch.cat(image_tokens, dim=0)
+                # ====================== tokenize images ======================
+
+
+                # ====================== tokenize text prompts ======================
+                if args.text_encoder_architecture in ("t5_clip", "gemma"):
+                    prompt_input_ids_clip = batch["prompt_input_ids"][0].to(accelerator.device, non_blocking=True)
+                    prompt_input_ids_t5 = batch["prompt_input_ids"][1].to(accelerator.device, non_blocking=True)
+                    prompt_input_ids_clip_1, prompt_input_ids_clip_2 = prompt_input_ids_clip.chunk(2, dim=0)
+                    prompt_input_ids_t5_1, prompt_input_ids_t5_2 = prompt_input_ids_t5.chunk(2, dim=0)
+                    encoder_hidden_states, cond_embeds = encode_prompt(
+                        [text_encoder, text_encoder_2], 
+                        [prompt_input_ids_clip_1, prompt_input_ids_t5_1], 
+                        args.text_encoder_architecture
+                    )
+                else:
+                    prompt_input_ids = batch["prompt_input_ids"].to(accelerator.device, non_blocking=True)
+                    prompt_input_ids_clip_1, prompt_input_ids_clip_2 = prompt_input_ids.chunk(2, dim=0)
+                    encoder_hidden_states, cond_embeds = encode_prompt(
+                        text_encoder, 
+                        prompt_input_ids_clip_1, 
+                        args.text_encoder_architecture
+                    )
+                encoder_hidden_states = encoder_hidden_states.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                cond_embeds = cond_embeds.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                # ====================== tokenize text prompts ======================
+
+
+                # ====================== image perturbation ======================
+                image_tokens_1, image_tokens_2 = image_tokens.chunk(2, dim=0)   # (b // 2, seq_len)
+                half_batch_size, seq_len = image_tokens_1.shape
+                sigma = torch.rand(half_batch_size, device=image_tokens_1.device)
+                image_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                image_mask_prob = image_mask_prob.clip(args.min_masking_rate)
+
+                num_token_masked = (seq_len * image_mask_prob).round().clamp(min=1)
+                batch_randperm = torch.rand(half_batch_size, seq_len, device=image_tokens_1.device).argsort(dim=-1)
+                mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                mask_id = accelerator.unwrap_model(model).config.vocab_size - 1
+                masked_image_ids = torch.where(mask, mask_id, image_tokens_1)
+                image_labels = torch.where(mask, image_tokens_1, -100)
+                # ====================== image perturbation ======================
+
+
+                # ====================== text perturbation  ======================
+                if args.text_encoder_architecture in ("t5_clip", "gemma"):
+                    half_batch_size, seq_len = prompt_input_ids_t5_2.shape
+                    sigma = torch.rand(half_batch_size, device=image_tokens_1.device)
+                    text_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                    text_mask_prob = text_mask_prob.clip(args.min_masking_rate)
+                    text_timestep = text_mask_prob.clone().clamp(min=1e-3)
+
+                    num_token_masked = (seq_len * text_mask_prob).round().clamp(min=1)
+                    batch_randperm = torch.rand(half_batch_size, seq_len, device=image_tokens_1.device).argsort(dim=-1)
+                    mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                    masked_prompt_input_ids_t5 = torch.where(mask, mask_id_2, prompt_input_ids_t5_2)
+                    text_labels = torch.where(mask, prompt_input_ids_t5_2, -100)
+                    
+                    # prepare input_ids for clip model
+                    batch_prompt_2 = []
+                    for i in range(masked_prompt_input_ids_t5.size(0)):
+                        masked_prompt_input_id = masked_prompt_input_ids_t5[i].tolist()
+                        prompt_2 = tokenizer_2.decode(masked_prompt_input_id, skip_special_tokens=True)
+                        batch_prompt_2.append(prompt_2)
+                    
+                    masked_prompt_input_ids_clip = tokenizer(
+                        batch_prompt_2,
+                        truncation=True,
+                        padding="max_length",
+                        max_length=77,
+                        return_tensors="pt"
+                    ).input_ids
+                    masked_prompt_input_ids_clip = masked_prompt_input_ids_clip.to(accelerator.device)
+                else:
+                    half_batch_size, seq_len = prompt_input_ids_clip_2.shape
+                    sigma = torch.rand(half_batch_size, device=image_tokens_1.device)
+                    text_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                    text_mask_prob = text_mask_prob.clip(args.min_masking_rate)
+                    text_timestep = text_mask_prob.clone().clamp(min=1e-3)
+
+                    num_token_masked = (seq_len * text_mask_prob).round().clamp(min=1)
+                    batch_randperm = torch.rand(half_batch_size, seq_len, device=image_tokens_1.device).argsort(dim=-1)
+                    mask = batch_randperm < num_token_masked.unsqueeze(-1)
+                    
+                    masked_prompt_input_ids_clip = torch.where(mask, mask_id_1, prompt_input_ids_clip_2)
+                    text_labels = torch.where(mask, prompt_input_ids_clip_2, -100)        
+                # ====================== text perturbation ======================
+
+
+                # ====================== encode masked text prompts ======================
+                if args.text_encoder_architecture in ("t5_clip", "gemma"):
+                    masked_encoder_hidden_states, masked_cond_embeds = encode_prompt(
+                        [text_encoder, text_encoder_2], 
+                        [masked_prompt_input_ids_clip, masked_prompt_input_ids_t5], 
+                        args.text_encoder_architecture
+                    )
+                else:
+                    masked_encoder_hidden_states, masked_cond_embeds = encode_prompt(
+                        text_encoder, 
+                        masked_prompt_input_ids_clip, 
+                        args.text_encoder_architecture
+                    )
+                masked_encoder_hidden_states = masked_encoder_hidden_states.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                masked_cond_embeds = masked_cond_embeds.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                # ====================== encode masked text prompts ======================
+                
+                
+                # for CFG
+                if args.cond_dropout_prob > 0.0:
+                    assert encoder_hidden_states is not None
+
+                    batch_size = encoder_hidden_states.shape[0]
+
+                    mask = (
+                        torch.zeros((batch_size, 1, 1), device=encoder_hidden_states.device).float().uniform_(0, 1)
+                        < args.cond_dropout_prob
+                    )
+
+                    empty_embeds_ = empty_embeds.expand(batch_size, -1, -1)
+                    encoder_hidden_states = torch.where(
+                        (encoder_hidden_states * mask).bool(), encoder_hidden_states, empty_embeds_
+                    )
+
+                    empty_clip_embeds_ = empty_clip_embeds.expand(batch_size, -1)
+                    cond_embeds = torch.where((cond_embeds * mask.squeeze(-1)).bool(), cond_embeds, empty_clip_embeds_)
+                    
+                vae_scale_factor = 2 ** (len(vq_model.config.block_out_channels) - 1)
+                resolution = args.resolution // vae_scale_factor
+                masked_image_ids = masked_image_ids.reshape(half_batch_size, resolution, resolution)
+                image_ids = image_tokens_2.reshape(half_batch_size, resolution, resolution)
+
+
+            # Train Step
+            with accelerator.accumulate(model):
+                codebook_size = accelerator.unwrap_model(model).config.codebook_size                   
+                if args.resolution == 1024: # only stage 3 and stage 4 do not apply 2*
+                    img_ids = _prepare_latent_image_ids(
+                        masked_image_ids.shape[0], 
+                        masked_image_ids.shape[-2],
+                        masked_image_ids.shape[-1],
+                        masked_image_ids.device,
+                        masked_image_ids.dtype
+                    )
+                else:
+                    img_ids = _prepare_latent_image_ids(
+                        masked_image_ids.shape[0],
+                        masked_image_ids.shape[-2],
+                        masked_image_ids.shape[-1],
+                        masked_image_ids.device,
+                        masked_image_ids.dtype
+                    )
+
+                txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(device=masked_image_ids.device, dtype=masked_image_ids.dtype)
+                
+                image_logits = (
+                    model(
+                        hidden_states=masked_image_ids, # should be (batch size, channel, height, width)
+                        encoder_hidden_states=encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                        micro_conds=image_micro_conds,
+                        pooled_projections=cond_embeds, # should be (batch_size, projection_dim)
+                        img_ids=img_ids,
+                        txt_ids=txt_ids,
+                        timestep=image_mask_prob,
+                    )[0]
+                    .reshape(half_batch_size, codebook_size, -1)
+                    .permute(0, 2, 1)
+                    .reshape(-1, codebook_size)
+                )
+
+                image_loss = F.cross_entropy(
+                    image_logits,
+                    image_labels.view(-1),
+                    ignore_index=-100,
+                    reduction="mean",
+                )
+
+                text_logits = model(
+                    hidden_states=image_ids, # should be (batch size, channel, height, width)
+                    encoder_hidden_states=masked_encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                    micro_conds=text_micro_conds,
+                    pooled_projections=masked_cond_embeds, # should be (batch_size, projection_dim)
+                    img_ids=img_ids,
+                    txt_ids=txt_ids,
+                    timestep=text_timestep,
+                )[1]
+                text_logits = text_logits.reshape(-1, accelerator.unwrap_model(model).config.tokenizer_vocab_size)
+
+                text_loss = F.cross_entropy(
+                    text_logits,
+                    text_labels.view(-1),
+                    ignore_index=-100,
+                    reduction="none",
+                )
+                text_loss = text_loss.reshape(half_batch_size, -1).mean(-1)
+                text_loss = text_loss / text_timestep
+                text_loss = text_loss.mean()
+
+                loss = image_loss + args.text_loss_weight * text_loss
+                
+                # Gather the losses across all processes for logging (if we use distributed training).
+                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                avg_masking_rate = accelerator.gather(text_mask_prob.repeat(args.train_batch_size)).mean()
+
+                accelerator.backward(loss)
+                
+                # Temporarily add this to identify unused parameters
+                # for name, param in accelerator.unwrap_model(model).named_parameters():
+                #     if param.grad is None:
+                #         print(f"Unused parameter: {name}")
+
+                if args.max_grad_norm is not None and accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                if (global_step + 1) % args.logging_steps == 0:
+                    logs = {
+                        "step_loss": avg_loss.item(),
+                        "lr": lr_scheduler.get_last_lr()[0],
+                        "avg_masking_rate": avg_masking_rate.item(),
+                    }
+                    accelerator.log(logs, step=global_step + 1)
+
+                    logger.info(
+                        f"Step: {global_step + 1} "
+                        f"Loss: {avg_loss.item():0.4f} "
+                        f"LR: {lr_scheduler.get_last_lr()[0]:0.6f}"
+                    )
+
+                if (global_step + 1) % args.checkpointing_steps == 0:
+                    save_checkpoint(args, accelerator, global_step + 1, logger)
+
+                if (global_step + 1) % args.validation_steps == 0 and accelerator.is_main_process:
+
+                    with torch.no_grad():
+                        logger.info("Generating images...")
+
+                        model.eval()
+
+                        scheduler = Scheduler.from_pretrained(
+                            args.pretrained_model_name_or_path,
+                            subfolder="scheduler",
+                            revision=args.revision,
+                            variant=args.variant,
+                        )
+                        
+                        pipe = UnifiedPipeline(
+                            transformer=accelerator.unwrap_model(model),
+                            tokenizer=tokenizer,
+                            text_encoder=text_encoder,
+                            vqvae=vq_model,
+                            scheduler=scheduler,
+                            tokenizer_2=tokenizer_2,
+                            text_encoder_2=text_encoder_2,
+                        )
+
+                        if not args.image_to_text_only:
+                            output = pipe(
+                                prompt=args.validation_prompts,
+                                height=args.resolution,
+                                width=args.resolution,
+                                guidance_scale=9,
+                                num_inference_steps=64,
+                            )
+                            pil_images = output.images
+
+                            result=[]
+                            for img in pil_images:
+                                if not isinstance(img, torch.Tensor):
+                                    img = transforms.ToTensor()(img)
+                                result.append(img.unsqueeze(0))
+                            result = torch.cat(result,dim=0)
+                            result = make_grid(result, nrow=3)
+                            save_image(result,os.path.join(args.output_dir, str(global_step)+'_text2image_1024_CFG-9.png'))
+                        
+                            output_data = {
+                                "step": global_step,
+                                "prompts": args.validation_prompts,
+                                "images": [f"{global_step}_text2image_1024_CFG-9_{i}.png" for i in range(len(pil_images))]
+                            }
+
+                            with open(os.path.join(args.output_dir, f"text2image_{global_step}.json"), "w") as f:
+                                json.dump(output_data, f, indent=2)
+
+                        image = load_images_to_tensor(args.validation_images, target_size=(args.resolution, args.resolution))
+                        output = pipe(
+                            prompt=args.validation_prompts, 
+                            height=args.resolution, 
+                            width=args.resolution, 
+                            guidance_scale=9,
+                            image=image,
+                            num_inference_steps=64
+                        )
+                        prompts = output.prompts
+                        
+                        output_data = {
+                            "step": global_step,
+                            "prompts": prompts,
+                        }
+
+                        with open(os.path.join(args.output_dir, f"image2text_{global_step}.json"), "w") as f:
+                            json.dump(output_data, f, indent=2)
+
+                        model.train()
+
+                global_step += 1
+
+            # Stop training if max steps is reached
+            if global_step >= args.max_train_steps:
+                break
+        # End for
+
+    accelerator.wait_for_everyone()
+
+    # Evaluate and save checkpoint at the end of training
+    save_checkpoint(args, accelerator, global_step, logger)
+
+    # Save the final trained checkpoint
+    if accelerator.is_main_process:
+        model = accelerator.unwrap_model(model)
+        model.save_pretrained(args.output_dir)
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main(parse_args())

train/train_unified.sh

@@ -0,0 +1,33 @@
+# bash it in root path
+PYTHON_PATH='./' accelerate launch --multi_gpu --gpu_ids '0,1,2,3,4,5,6,7' --main_process_port 25000 --num_processes 8 train/train_unified.py \
+        --output_dir "/path/to/output/dir" \
+        --train_batch_size 8 \
+        --gradient_accumulation_steps 8 \
+        --learning_rate 1e-4 \
+        --text_loss_weight 0 \
+        --max_grad_norm 10 \
+        --pretrained_model_name_or_path "MeissonFlow/Meissonic" \
+        --pretrained_transformer_path "MeissonFlow/Meissonic" \
+        --text_encoder_architecture 'open_clip' \
+        --instance_dataset 'ImageCaptionLargeDataset' \
+        --instance_data_dir  '/path/to/data/' \
+        --resolution 512 \
+        --mixed_precision fp16 \
+        --lr_scheduler constant \
+        --use_8bit_adam \
+        --dataloader_num_workers 4 \
+        --validation_prompts \
+            'a boy' \
+            'A serene mountain landscape with towering snow-capped peaks, a crystal-clear blue lake reflecting the mountains, dense pine forests, and a vibrant orange sunrise illuminating the sky.' \
+            'A playful golden retriever puppy with a shiny coat, bounding through a meadow filled with colorful wildflowers, under a bright, clear blue sky.' \
+            'A bustling city street at night, illuminated by vibrant neon signs in various colors, with busy pedestrians, street vendors, and a light rain creating reflective puddles on the pavement.' \
+            'A majestic, medieval castle perched on a rugged cliffside, overlooking a vast, calm ocean at sunset, with the sky painted in hues of pink, orange, and purple.' \
+            'An elegant ballerina in a white tutu, dancing gracefully on a grand stage with ornate, gold-trimmed curtains, under a spotlight that casts a soft glow.' \
+            'A cozy, rustic log cabin nestled in a snow-covered forest, with smoke rising from the stone chimney, warm lights glowing from the windows, and a path of footprints leading to the front door.'\
+            'A Cute Cat' \
+            'A Snow Mountain'\
+        --max_train_steps 100000 \
+        --checkpointing_steps 1000 \
+        --validation_steps 100 \
+        --report_to 'wandb' \
+        --logging_steps 10

train/train_unified_new.py

@@ -0,0 +1,1078 @@
+# 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 argparse
+import copy
+import logging
+import math
+import os
+from pathlib import Path
+import sys
+sys.path.append(os.getcwd())
+import json
+import gc
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from peft import LoraConfig
+from peft.utils import get_peft_model_state_dict
+from torch.utils.data import DataLoader
+from torchvision import transforms
+
+from transformers.models.gemma2.modeling_gemma2 import Gemma2Model
+from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
+from transformers import (
+    CLIPTextModelWithProjection,
+    CLIPTokenizer,
+    CLIPImageProcessor,
+    CLIPVisionModelWithProjection,
+)
+
+import diffusers.optimization
+from diffusers import VQModel
+
+from src.scheduler import Scheduler
+from diffusers.loaders import LoraLoaderMixin
+from diffusers.utils import is_wandb_available
+from src.pipeline import UnifiedPipeline_new
+from torchvision.utils import save_image, make_grid
+from train.trainer_utils import save_checkpoint
+from train.dataset_utils import ImageCaptionLargeDataset
+from train.dataset_utils import encode_prompt
+from src.transformer import SymmetricTransformer2DModel
+from train.trainer_utils import load_images_to_tensor
+
+if is_wandb_available():
+    import wandb
+    # wandb.login(key="")
+
+logger = get_logger(__name__, log_level="INFO")
+
+import torch._dynamo
+torch._dynamo.config.verbose = True
+
+# Optionally suppress errors to fall back to eager execution
+torch._dynamo.config.suppress_errors = True
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--pretrained_transformer_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--text_encoder_architecture",
+        type=str,
+        default="open_clip",
+        required=False,
+        help="The architecture of the text encoder. One of ['open_clip', 'gemma']",
+    )
+    parser.add_argument(
+        "--clip_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--text_encoder_2_name_or_path",
+        type=str,
+        default=None,
+        required=False,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--instance_dataset",
+        type=str,
+        default=None,
+        required=False,
+        help="The dataset to use for training. One of ['MSCOCO600K', 'PickaPicV2']",
+    )
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--training_from_scratch",
+        type=bool,
+        default=False,
+        required=False
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--variant",
+        type=str,
+        default=None,
+        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+    )
+    parser.add_argument(
+        "--instance_data_image", type=str, default=None, required=False, help="A single training image"
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.9999)
+    parser.add_argument("--ema_update_after_step", type=int, default=0)
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="muse_training",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
+            "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
+            "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
+            "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
+            "instructions."
+        ),
+    )
+    parser.add_argument(
+        "--logging_steps",
+        type=int,
+        default=50,
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=(
+            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
+            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
+            " for more details"
+        ),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--text_loss_reduction",
+        type=str,
+        default="mean",
+        help="The reduction method for the text loss. One of ['mean', 'reweighted']",
+    )
+    parser.add_argument(
+        "--text_loss_weight",
+        type=float,
+        default=0.2,
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=0.0003,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--validation_steps",
+        type=int,
+        default=100,
+        help=(
+            "Run validation every X steps. Validation consists of running the prompt"
+            " `args.validation_prompt` multiple times: `args.num_validation_images`"
+            " and logging the images."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="wandb",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument("--validation_prompts", type=str, nargs="*")
+    parser.add_argument("--validation_images", type=str, default="./assets")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument("--split_vae_encode", type=int, required=False, default=None)
+    parser.add_argument("--min_masking_rate", type=float, default=0.0)
+    parser.add_argument("--cond_dropout_prob", type=float, default=0.0)
+    parser.add_argument("--max_grad_norm", default=50.0, type=float, help="Max gradient norm.", required=False)
+    parser.add_argument("--use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--text_encoder_use_lora", action="store_true", help="Fine tune the model using LoRa")
+    parser.add_argument("--lora_r", default=16, type=int)
+    parser.add_argument("--lora_alpha", default=32, type=int)
+    parser.add_argument("--lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--text_encoder_lora_r", default=16, type=int)
+    parser.add_argument("--text_encoder_lora_alpha", default=32, type=int)
+    parser.add_argument("--text_encoder_lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
+    parser.add_argument("--train_text_encoder", action="store_true")
+    parser.add_argument("--image_to_text_only", action="store_true")
+    parser.add_argument("--image_key", type=str, required=False)
+    parser.add_argument("--prompt_key", type=str, required=False)
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--prompt_prefix", type=str, required=False, default=None)
+
+    args = parser.parse_args()
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+    num_datasources = sum(
+        [x is not None for x in [args.instance_data_dir, args.instance_data_image]]
+    )
+
+    if num_datasources != 1:
+        raise ValueError(
+            "provide one and only one of `--instance_data_dir`, `--instance_data_image`, or `--instance_data_dataset`"
+        )
+
+    if args.instance_data_dir is not None:
+        if not os.path.exists(args.instance_data_dir):
+            raise ValueError(f"Does not exist: `--args.instance_data_dir` {args.instance_data_dir}")
+
+    if args.instance_data_image is not None:
+        if not os.path.exists(args.instance_data_image):
+            raise ValueError(f"Does not exist: `--args.instance_data_image` {args.instance_data_image}")
+
+    return args
+
+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 main(args):
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if accelerator.is_main_process:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+
+    # if accelerator.is_main_process:
+    #     accelerator.init_trackers("meissonic", config=vars(copy.deepcopy(args)))
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Initialize image processor and image encoder (CLIP Vision Model with Projection)
+    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+        args.clip_model_name_or_path
+    )
+    image_processor = CLIPImageProcessor.from_pretrained(
+        args.clip_model_name_or_path
+    )
+
+    # Freeze image encoder parameters
+    image_encoder.requires_grad_(False)
+
+    text_encoder = CLIPTextModelWithProjection.from_pretrained(
+        args.clip_model_name_or_path
+    )
+    tokenizer = CLIPTokenizer.from_pretrained(
+        args.clip_model_name_or_path
+    )
+
+    text_encoder.requires_grad_(False)
+
+    if args.text_encoder_architecture == "open_clip":
+        tokenizer_2 = None
+        text_encoder_2 = None
+
+        mask_token = "<mask>"
+        num_new_tokens = tokenizer.add_tokens(mask_token)
+        mask_id_1 = tokenizer.convert_tokens_to_ids(mask_token)
+        if num_new_tokens > 0:
+            text_encoder.resize_token_embeddings(len(tokenizer))
+            mask_token_embedding = text_encoder.get_input_embeddings().weight[mask_id_1]
+            mask_token_embedding = mask_token_embedding.clone().detach().cpu().float()
+            if accelerator.is_main_process:
+                print("Saving masked token embedding...")
+                torch.save(mask_token_embedding, os.path.join(args.output_dir, "mask_token_embedding.pth"))
+                            
+    elif args.text_encoder_architecture == "gemma":        
+        tokenizer_2 = GemmaTokenizerFast.from_pretrained(
+            args.text_encoder_name_or_path,
+        )
+        text_encoder_2 = Gemma2Model.from_pretrained(
+            args.text_encoder_name_or_path,
+        )
+
+        mask_token = "<mask>"
+        
+        num_new_tokens = tokenizer_2.add_tokens(mask_token)
+        mask_id_2 = tokenizer_2.convert_tokens_to_ids(mask_token)
+
+        if num_new_tokens > 0:
+            text_encoder_2.resize_token_embeddings(len(tokenizer_2))
+            mask_token_embedding = text_encoder_2.get_input_embeddings().weight[mask_id_2]
+            mask_token_embedding = mask_token_embedding.clone().detach().cpu().float()
+            if accelerator.is_main_process:
+                print("Saving masked token embedding...")
+                torch.save(mask_token_embedding, os.path.join(args.output_dir, "mask_token_embedding.pth"))
+                            
+        text_encoder_2.requires_grad_(False)
+    else:
+        raise ValueError(f"Unknown text encoder architecture: {args.text_encoder_architecture}")
+    
+    vq_model = VQModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vqvae", revision=args.revision, variant=args.variant
+    )
+    vq_model.requires_grad_(False)
+
+    model = SymmetricTransformer2DModel.from_pretrained(
+        args.pretrained_transformer_path, 
+        subfolder="transformer", 
+        low_cpu_mem_usage=False,
+        device_map=None
+    )
+
+    if model.config.tokenizer_vocab_size is None:
+        if args.text_encoder_architecture == "open_clip":
+            model.register_to_config(tokenizer_vocab_size=len(tokenizer))
+        elif args.text_encoder_architecture == "gemma":
+            model.register_to_config(tokenizer_vocab_size=len(tokenizer_2))
+        else:
+            raise ValueError(f"Unknown text encoder architecture!")
+
+        if accelerator.is_main_process:
+            print(f"model's tokenizer vocab size is {model.config.tokenizer_vocab_size}")
+            
+        model.text_decoder = nn.Sequential(
+            nn.LayerNorm(model.inner_dim, elementwise_affine=False, eps=1e-6),
+            nn.Linear(model.inner_dim, model.config.tokenizer_vocab_size, bias=False)
+        )
+
+    model = torch.compile(model)
+
+    if args.use_lora:
+        lora_config = LoraConfig(
+            r=args.lora_r,
+            lora_alpha=args.lora_alpha,
+            target_modules=args.lora_target_modules,
+        )
+        model.add_adapter(lora_config)
+
+    model.train()
+
+    if args.image_to_text_only:
+        frozen_keys = ["project_from_hidden", "up_block", "mlm_layer"]
+        for n, p in model.named_parameters():
+            if any([frozen_key in n for frozen_key in frozen_keys]):
+                p.requires_grad_(False)
+            else:
+                p.requires_grad_(True)
+    else:
+        model.requires_grad_(True)
+
+    if args.gradient_checkpointing:
+        model.enable_gradient_checkpointing()   
+
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            transformer_lora_layers_to_save = None
+            text_encoder_lora_layers_to_save = None
+
+            for model_ in models:
+                if isinstance(model_, type(accelerator.unwrap_model(model))):
+                    if args.use_lora:
+                        transformer_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "transformer"))
+                elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                    if args.text_encoder_use_lora:
+                        text_encoder_lora_layers_to_save = get_peft_model_state_dict(model_)
+                    else:
+                        model_.save_pretrained(os.path.join(output_dir, "text_encoder"))
+                else:
+                    raise ValueError(f"unexpected save model: {model_.__class__}")
+
+                # make sure to pop weight so that corresponding model is not saved again
+                weights.pop()
+
+            if transformer_lora_layers_to_save is not None or text_encoder_lora_layers_to_save is not None:
+                LoraLoaderMixin.save_lora_weights(
+                    output_dir,
+                    unet_lora_layers=transformer_lora_layers_to_save,
+                    text_encoder_lora_layers=text_encoder_lora_layers_to_save,
+                )
+                
+
+    def load_model_hook(models, input_dir):
+        transformer = None
+        text_encoder_ = None
+
+        # this part is added for keep consistency when add model.compile() in the model
+        def adap_compile(ori_dict):#add '_orig_mod.' to each key
+            new_dict = {}
+            for k,v in ori_dict.items():
+                new_dict['_orig_mod.' + k] = v
+            return new_dict
+            
+        while len(models) > 0:
+            model_ = models.pop()
+
+            if isinstance(model_, type(accelerator.unwrap_model(model))):
+                if args.use_lora:
+                    transformer = model_
+                else:
+                    load_model = SymmetricTransformer2DModel.from_pretrained(os.path.join(input_dir, "transformer"), low_cpu_mem_usage=False, device_map=None)
+                    model_.load_state_dict(adap_compile(load_model.state_dict()))
+                    del load_model
+            elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
+                if args.text_encoder_use_lora:
+                    text_encoder_ = model_
+                else:
+                    try:
+                        load_model = CLIPTextModelWithProjection.from_pretrained(os.path.join(input_dir, "text_encoder"))
+                        model_.load_state_dict(load_model.state_dict())
+                        # print('finished loading text encoder!')
+                    except:
+                        print('Not found text-encoder model in current folder. So we download one text encoder from Internet.')
+                        load_model = CLIPTextModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
+                        model_.load_state_dict(load_model.state_dict())
+                    del load_model
+            else:
+                raise ValueError(f"unexpected save model: {model.__class__}")
+
+        if transformer is not None or text_encoder_ is not None:
+            lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir)
+            LoraLoaderMixin.load_lora_into_text_encoder(
+                lora_state_dict, network_alphas=network_alphas, text_encoder=text_encoder_
+            )
+            LoraLoaderMixin.load_lora_into_transformer(
+                lora_state_dict, network_alphas=network_alphas, transformer=transformer
+            )
+
+    accelerator.register_load_state_pre_hook(load_model_hook)
+    accelerator.register_save_state_pre_hook(save_model_hook)
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+        )
+
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+            )
+
+        optimizer_cls = bnb.optim.AdamW8bit
+    else:
+        optimizer_cls = torch.optim.AdamW
+
+    optimizer_grouped_parameters = [
+        {
+            "params": [p for p in model.parameters() if p.requires_grad],
+            "weight_decay": args.adam_weight_decay,
+        }
+    ]
+    optimizer = optimizer_cls(
+        optimizer_grouped_parameters,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    logger.info("Creating dataloaders and lr_scheduler")
+
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+
+    if args.text_encoder_architecture in ("t5_clip", "gemma"):
+        tokenizer_for_dataset = [tokenizer, tokenizer_2]
+    else:
+        tokenizer_for_dataset = tokenizer
+
+    if args.instance_dataset == "ImageCaptionLargeDataset":
+        dataset = ImageCaptionLargeDataset(
+            root_dir=args.instance_data_dir,
+            tokenizer=tokenizer_for_dataset,
+            size=args.resolution,
+            text_encoder_architecture=args.text_encoder_architecture
+        )
+    elif args.instance_dataset == "DATA_TYPE":
+        raise NotImplementedError("DATA_TYPE is not yet supported")
+    else:
+        assert False
+
+    def collate_fn(samples):
+        images = [sample["image"] for sample in samples]
+        micro_conds = [sample["micro_conds"] for sample in samples]
+        
+        images = torch.stack(images, dim=0)
+        micro_conds = torch.stack(micro_conds, dim=0)
+        
+        if isinstance(samples[0]["prompt_input_ids"], list):
+            input_ids = [sample["prompt_input_ids"][0] for sample in samples]
+            input_ids_2 = [sample["prompt_input_ids"][1] for sample in samples]
+            
+            input_ids = torch.cat(input_ids, dim=0)
+            input_ids_2 = torch.cat(input_ids_2, dim=0)
+            prompt_input_ids = [input_ids, input_ids_2]
+        else:
+            input_ids = [sample["prompt_input_ids"] for sample in samples]
+            
+            input_ids = torch.cat(input_ids, dim=0)
+            prompt_input_ids = input_ids
+        
+        ret = dict(
+            images=images,
+            micro_conds=micro_conds,
+            prompt_input_ids=prompt_input_ids,
+        )
+        
+        return ret
+    
+    train_dataloader = DataLoader(
+        dataset,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        num_workers=args.dataloader_num_workers,
+        collate_fn=collate_fn,
+        pin_memory=True,
+    )
+    train_dataloader.num_batches = len(train_dataloader)
+
+    lr_scheduler = diffusers.optimization.get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+    )
+
+    logger.info("Preparing model, optimizer and dataloaders")
+
+    model, optimizer, lr_scheduler, train_dataloader = accelerator.prepare(
+        model, optimizer, lr_scheduler, train_dataloader
+    )
+
+    train_dataloader.num_batches = len(train_dataloader)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    image_encoder.to(device=accelerator.device, dtype=weight_dtype)
+
+    if args.text_encoder_architecture in ("t5_clip", "gemma"):
+        text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+        text_encoder_2.to(device=accelerator.device, dtype=weight_dtype)
+    else:
+        text_encoder.to(device=accelerator.device, dtype=weight_dtype)
+            
+    vq_model.to(device=accelerator.device)
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(train_dataloader.num_batches / args.gradient_accumulation_steps)
+    # Afterwards we recalculate our number of training epochs.
+    # Note: We are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Train!
+    logger.info("***** Running training *****")
+    logger.info(f"  Num training steps = {args.max_train_steps}")
+    logger.info(f"  Instantaneous batch size per device = { args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+
+    resume_from_checkpoint = args.resume_from_checkpoint
+    if resume_from_checkpoint:
+        if resume_from_checkpoint == "latest":
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            if len(dirs) > 0:
+                resume_from_checkpoint = os.path.join(args.output_dir, dirs[-1])
+            else:
+                resume_from_checkpoint = None
+
+        if resume_from_checkpoint is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+        else:
+            accelerator.print(f"Resuming from checkpoint {resume_from_checkpoint}")
+
+    if resume_from_checkpoint is None:
+        global_step = 0
+        first_epoch = 0
+    else:
+        accelerator.load_state(resume_from_checkpoint)
+        global_step = int(os.path.basename(resume_from_checkpoint).split("-")[1])
+        first_epoch = global_step // num_update_steps_per_epoch
+
+    # As stated above, we are not doing epoch based training here, but just using this for book keeping and being able to
+    # reuse the same training loop with other datasets/loaders.
+    for epoch in range(first_epoch, num_train_epochs):
+        for batch in train_dataloader:
+            torch.cuda.empty_cache()
+            with torch.no_grad():
+                micro_conds = batch["micro_conds"].to(accelerator.device, non_blocking=True)
+                gen_micro_conds, und_micro_conds = micro_conds.chunk(2, dim=0)
+                
+                pixel_values = batch["images"].to(accelerator.device, non_blocking=True)    # [b, 3, res, res]
+                batch_size = pixel_values.shape[0]
+                half_batch_size = batch_size // 2
+
+                # ====================== tokenize images ======================
+                image_tokens = vq_model.quantize(
+                    vq_model.encode(pixel_values).latents
+                )[2][2].reshape(batch_size, -1)  # [b, seq_len]
+                # ====================== tokenize images ======================
+
+
+                # ====================== tokenize text prompts ======================
+                if args.text_encoder_architecture == "gemma":
+                    prompt_input_ids_1 = batch["prompt_input_ids"][0].to(accelerator.device, non_blocking=True)
+                    prompt_input_ids_2 = batch["prompt_input_ids"][1].to(accelerator.device, non_blocking=True)
+
+                    prompt_input_ids_gen_1, prompt_input_ids_und_1 = prompt_input_ids_1.chunk(2, dim=0)
+                    prompt_input_ids_gen_2, prompt_input_ids_und_2 = prompt_input_ids_2.chunk(2, dim=0)
+                else:
+                    prompt_input_ids = batch["prompt_input_ids"].to(accelerator.device, non_blocking=True)
+                    prompt_input_ids_gen, prompt_input_ids_und = prompt_input_ids.chunk(2, dim=0)
+                # ====================== tokenize text prompts ======================
+
+
+                # ====================== image perturbation ======================
+                image_tokens_gen, image_tokens_und = image_tokens.chunk(2, dim=0)   # (b // 2, seq_len)
+                _, seq_len = image_tokens_gen.shape
+                sigma = torch.rand(half_batch_size, device=image_tokens_gen.device)
+                gen_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                gen_mask_prob = gen_mask_prob.clip(args.min_masking_rate)
+
+                num_token_masked = (seq_len * gen_mask_prob).round().clamp(min=1)
+                batch_randperm = torch.rand(half_batch_size, seq_len, device=image_tokens_gen.device).argsort(dim=-1)
+                mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                mask_id = accelerator.unwrap_model(model).config.vocab_size - 1
+                masked_image_ids = torch.where(mask, mask_id, image_tokens_gen)
+                image_labels = torch.where(mask, image_tokens_gen, -100)
+                # ====================== image perturbation ======================
+
+
+                # ====================== text perturbation  ======================
+                if args.text_encoder_architecture in ("t5_clip", "gemma"):
+                    half_batch_size, seq_len = prompt_input_ids_und_2.shape
+                    sigma = torch.rand(half_batch_size, device=image_tokens_gen.device)
+                    text_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                    text_mask_prob = text_mask_prob.clip(args.min_masking_rate)
+                    text_timestep = text_mask_prob.clone()
+
+                    num_token_masked = (seq_len * text_mask_prob).round().clamp(min=1)
+                    batch_randperm = torch.rand(half_batch_size, seq_len, device=image_tokens_gen.device).argsort(dim=-1)
+                    mask = batch_randperm < num_token_masked.unsqueeze(-1)
+
+                    masked_prompt_input_ids_und = torch.where(mask, mask_id_2, prompt_input_ids_und_2)
+                    text_labels = torch.where(mask, prompt_input_ids_und_2, -100)
+                    
+                else:
+                    half_batch_size, seq_len = prompt_input_ids_und.shape
+                    sigma = torch.rand(half_batch_size, device=image_tokens_gen.device)
+                    text_mask_prob = torch.cos(sigma * math.pi * 0.5)
+                    text_mask_prob = text_mask_prob.clip(args.min_masking_rate)
+                    text_timestep = text_mask_prob.clone().clamp(min=1e-3)
+
+                    num_token_masked = (seq_len * text_mask_prob).round().clamp(min=1)
+                    batch_randperm = torch.rand(half_batch_size, seq_len, device=image_tokens_gen.device).argsort(dim=-1)
+                    mask = batch_randperm < num_token_masked.unsqueeze(-1)
+                    
+                    masked_prompt_input_ids_und = torch.where(mask, mask_id_1, prompt_input_ids_und)
+                    text_labels = torch.where(mask, prompt_input_ids_und, -100)        
+                # ====================== text perturbation ======================
+
+
+                # ====================== encode text prompts ======================
+                if args.text_encoder_architecture == "gemma":
+                    masked_encoder_hidden_states, _ = encode_prompt(
+                        [text_encoder, text_encoder_2],
+                        [prompt_input_ids_und_1, masked_prompt_input_ids_und], 
+                        args.text_encoder_architecture
+                    )
+                    encoder_hidden_states, text_pooled_embeds = encode_prompt(
+                        [text_encoder, text_encoder_2], 
+                        [prompt_input_ids_gen_1, prompt_input_ids_gen_2], 
+                        args.text_encoder_architecture
+                    )
+                else:
+                    masked_encoder_hidden_states, _ = encode_prompt(
+                        text_encoder, 
+                        masked_prompt_input_ids_und, 
+                        args.text_encoder_architecture
+                    )
+                    encoder_hidden_states, text_pooled_embeds = encode_prompt(
+                        text_encoder, 
+                        prompt_input_ids_gen, 
+                        args.text_encoder_architecture
+                    )
+                # obtain the cond_embeds through send pixel_values[half_batch_size:] to image_encoder, we use the clip pooled embedding as cond_embeds
+                # pixel_values need image process, the value of pixel values arange from 0 to 1
+                with torch.no_grad():
+                    processed_pixel_values = image_processor(
+                        pixel_values[half_batch_size:],
+                        do_rescale=False,
+                        do_resize=True,
+                        do_normalize=True,
+                        return_tensors="pt"
+                    )["pixel_values"].to(image_encoder.device, dtype=image_encoder.dtype)
+                    image_pooled_embeds = image_encoder(processed_pixel_values).image_embeds
+
+                # for text-to-image
+                encoder_hidden_states = encoder_hidden_states.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                text_pooled_embeds = text_pooled_embeds.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                # for image-to-text
+                masked_encoder_hidden_states = masked_encoder_hidden_states.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                image_pooled_embeds = image_pooled_embeds.to(accelerator.device, dtype=accelerator.unwrap_model(model).dtype)
+                # ====================== encode text prompts ======================
+                    
+                vae_scale_factor = 2 ** (len(vq_model.config.block_out_channels) - 1)
+                resolution = args.resolution // vae_scale_factor
+                masked_image_ids = masked_image_ids.reshape(half_batch_size, resolution, resolution)
+                image_ids = image_tokens_und.reshape(half_batch_size, resolution, resolution)
+
+            # Train Step
+            with accelerator.accumulate(model):
+                codebook_size = accelerator.unwrap_model(model).config.codebook_size                   
+                img_ids = _prepare_latent_image_ids(
+                    masked_image_ids.shape[0], 
+                    masked_image_ids.shape[-2],
+                    masked_image_ids.shape[-1],
+                    masked_image_ids.device,
+                    masked_image_ids.dtype
+                )
+
+                txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(device=masked_image_ids.device, dtype=masked_image_ids.dtype)
+                
+                image_logits = (
+                    model(
+                        hidden_states=masked_image_ids, # should be (batch size, channel, height, width)
+                        encoder_hidden_states=encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                        micro_conds=gen_micro_conds,
+                        pooled_projections=text_pooled_embeds, # should be (batch_size, projection_dim)
+                        img_ids=img_ids,
+                        txt_ids=txt_ids,
+                        timestep=gen_mask_prob,
+                    )[0]
+                    .reshape(half_batch_size, codebook_size, -1)
+                    .permute(0, 2, 1)
+                    .reshape(-1, codebook_size)
+                )
+
+                image_loss = F.cross_entropy(
+                    image_logits,
+                    image_labels.view(-1),
+                    ignore_index=-100,
+                    reduction="mean",
+                )
+
+                text_logits = model(
+                    hidden_states=image_ids, # should be (batch size, channel, height, width)
+                    encoder_hidden_states=masked_encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
+                    micro_conds=und_micro_conds,
+                    pooled_projections=image_pooled_embeds, # should be (batch_size, projection_dim)
+                    img_ids=img_ids,
+                    txt_ids=txt_ids,
+                    timestep=text_mask_prob,
+                )[1]
+                text_logits = text_logits.reshape(-1, accelerator.unwrap_model(model).config.tokenizer_vocab_size)
+
+                if args.text_loss_reduction == "mean":
+                    text_loss = F.cross_entropy(
+                        text_logits,
+                        text_labels.view(-1),
+                        ignore_index=-100,
+                        reduction="mean",
+                    )
+                elif args.text_loss_reduction == "reweighted":
+                    text_loss = F.cross_entropy(
+                        text_logits,
+                        text_labels.view(-1),
+                        ignore_index=-100,
+                        reduction="none",
+                    )
+                    text_loss = text_loss.reshape(half_batch_size, -1).mean(-1)
+                    text_loss = text_loss / text_timestep
+                    text_loss = text_loss.mean()
+                else:
+                    raise ValueError(f"Unknown text_loss_reduction: {args.text_loss_reduction}")
+
+                loss = image_loss + args.text_loss_weight * text_loss
+                
+                # Gather the losses across all processes for logging (if we use distributed training).
+                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+                avg_masking_rate = accelerator.gather(gen_mask_prob.repeat(args.train_batch_size)).mean()
+
+                accelerator.backward(loss)
+
+                if args.max_grad_norm is not None and accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+                optimizer.step()
+                lr_scheduler.step()
+
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                if (global_step + 1) % args.logging_steps == 0:
+                    logs = {
+                        "step_loss": avg_loss.item(),
+                        "lr": lr_scheduler.get_last_lr()[0],
+                        "avg_masking_rate": avg_masking_rate.item(),
+                    }
+                    accelerator.log(logs, step=global_step + 1)
+
+                    logger.info(
+                        f"Step: {global_step + 1} "
+                        f"Loss: {avg_loss.item():0.4f} "
+                        f"LR: {lr_scheduler.get_last_lr()[0]:0.6f}"
+                    )
+
+                if (global_step + 1) % args.checkpointing_steps == 0:
+                    save_checkpoint(args, accelerator, global_step + 1, logger)
+
+                if (global_step + 1) % args.validation_steps == 0 and accelerator.is_main_process:
+
+                    with torch.no_grad():
+                        logger.info("Evaluating...")
+
+                        model.eval()
+
+                        scheduler = Scheduler.from_pretrained(
+                            args.pretrained_model_name_or_path,
+                            subfolder="scheduler",
+                            revision=args.revision,
+                            variant=args.variant,
+                        )
+                        
+                        pipe = UnifiedPipeline_new(
+                            transformer=accelerator.unwrap_model(model),
+                            tokenizer=tokenizer,
+                            text_encoder=text_encoder,
+                            vqvae=vq_model,
+                            scheduler=scheduler,
+                            tokenizer_2=tokenizer_2,
+                            text_encoder_2=text_encoder_2,
+                            clip_image_processor=image_processor,
+                            image_encoder=image_encoder,
+                        )
+
+                        if not args.image_to_text_only:
+                            output = pipe(
+                                prompt=args.validation_prompts,
+                                height=args.resolution,
+                                width=args.resolution,
+                                guidance_scale=9,
+                                num_inference_steps=64,
+                            )
+                            pil_images = output.images
+
+                            result=[]
+                            for img in pil_images:
+                                if not isinstance(img, torch.Tensor):
+                                    img = transforms.ToTensor()(img)
+                                result.append(img.unsqueeze(0))
+                            result = torch.cat(result,dim=0)
+                            result = make_grid(result, nrow=3)
+                            save_image(result,os.path.join(args.output_dir, str(global_step)+'_text2image_1024_CFG-9.png'))
+
+                        image = load_images_to_tensor(args.validation_images, target_size=(args.resolution, args.resolution))
+                        output = pipe(
+                            height=args.resolution, 
+                            width=args.resolution, 
+                            guidance_scale=9,
+                            image=image,
+                            num_inference_steps=64
+                        )
+                        prompts = output.prompts
+                        
+                        output_data = {
+                            "step": global_step,
+                            "prompts": prompts,
+                        }
+
+                        with open(os.path.join(args.output_dir, f"image2text_{global_step}.json"), "w") as f:
+                            json.dump(output_data, f, indent=2)
+
+                        model.train()
+
+                global_step += 1
+
+            # Stop training if max steps is reached
+            if global_step >= args.max_train_steps:
+                break
+        # End for
+
+    accelerator.wait_for_everyone()
+
+    # Evaluate and save checkpoint at the end of training
+    save_checkpoint(args, accelerator, global_step, logger)
+
+    # Save the final trained checkpoint
+    if accelerator.is_main_process:
+        model = accelerator.unwrap_model(model)
+        model.save_pretrained(args.output_dir)
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    main(parse_args())

train/trainer_utils.py

@@ -0,0 +1,101 @@
+# 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 shutil
+from pathlib import Path, PosixPath
+
+import torch
+from PIL import Image
+from torchvision import transforms
+
+
+def save_checkpoint(args, accelerator, global_step, logger):
+    output_dir = args.output_dir
+
+    # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+    if accelerator.is_main_process and args.checkpoints_total_limit is not None:
+        checkpoints = os.listdir(output_dir)
+        checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+        checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+        # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+        if len(checkpoints) >= args.checkpoints_total_limit:
+            num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+            removing_checkpoints = checkpoints[0:num_to_remove]
+
+            logger.info(
+                f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+            )
+            logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+            for removing_checkpoint in removing_checkpoints:
+                removing_checkpoint = os.path.join(output_dir, removing_checkpoint)
+                shutil.rmtree(removing_checkpoint)
+
+    save_path = Path(output_dir) / f"checkpoint-{global_step}"
+    accelerator.save_state(save_path)
+    logger.info(f"Saved state to {save_path}")
+
+
+def load_images_to_tensor(path, target_size=(1024, 1024)):
+    """
+    Args:
+        folder_path
+        target_size: (height, width)
+    
+    Return:
+        torch.Tensor: [B, 3, H, W] in [0, 1]
+    """
+    valid_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp')
+    
+    if isinstance(path, list):
+        image_files = path
+    elif isinstance(path, str) and os.path.isdir(path):
+        image_files = [f for f in os.listdir(path) if f.lower().endswith(valid_extensions)]
+    elif isinstance(path, str):
+        image_files = [path]
+    else:
+        raise ValueError(f"Unsupported folder_path type: {type(path)}")
+    
+    if not image_files:
+        raise ValueError(f"No valid images found in {path}")
+    
+    transform = transforms.Compose([
+        transforms.Resize(target_size),
+        transforms.ToTensor(),
+    ])
+    
+    tensors = []
+    for img_file in image_files:
+        try:
+            if isinstance(path, str) and os.path.isdir(path):
+                img_path = os.path.join(path, img_file)
+            else:
+                img_path = img_file
+            img = Image.open(img_path).convert('RGB')
+            tensor = transform(img)
+            tensors.append(tensor)
+        except Exception as e:
+            print(f"Error processing {img_file}: {e}")
+    
+    if not tensors:
+        raise ValueError("No images could be loaded")
+    
+    batch_tensor = torch.stack(tensors)
+    
+    assert batch_tensor.shape[1:] == (3, *target_size), \
+        f"Output shape is {batch_tensor.shape}, expected (B, 3, {target_size[0]}, {target_size[1]})"
+    
+    return batch_tensor