Add project files

.gitignore

@@ -124,6 +124,17 @@ server.log
 # Distribution of plugin
 dist_upload/
 
+# Explicitly ignore built egg-info metadata
+comfyui_longcat.egg-info/
+
+# Pytest cache and artifacts
+.pytest_cache/
+
+# Explicit compiled python and cache files
+*.pyc
+__pycache__/
+*/__pycache__/
+
 # Other common
 *.swp
 *~

README.md

@@ -0,0 +1,53 @@
+# ComfyUI LongCat Plugin
+
+[简体中文](README_CN.md)
+
+## Overview
+ComfyUI nodes wrapping LongCat image generation and editing pipelines (diffusers-based). Includes text-to-image and multi-image edit flows with prompt rewriting, aspect-aware sizing, and latent decoding for ComfyUI.
+
+## Features
+- **LongCatCheckpointLoader**: Loads the LongCat model components (Model, CLIP, VAE).
+- **TextEncodeLongCatImage / TextEncodeLongCatImageEdit**: Encodes text prompts and reference images (for editing) into conditioning.
+- **VAEEncodeLongCat / VAEDecodeLongCat**: Encodes images to latents and decodes latents to images using the LongCat VAE.
+- **LongCatSampler**: Handles the sampling process (denoising) using the LongCat transformer.
+- **LongCatImageSizeScale**: Scales images to a target pixel area and rounds dimensions to multiples of 16.
+
+## Installation
+1.  **Copy Folder**: Copy the `comfyui_longcat` folder into your ComfyUI `custom_nodes` directory.
+    *   Example: `ComfyUI/custom_nodes/comfyui_longcat`
+2.  **Install Dependencies**: Ensure your ComfyUI environment has the required packages:
+    ```bash
+    pip install -r requirements.txt
+    ```
+    (Note: `requirements.txt` is in the root of this repo, or check `setup.py` for list).
+3.  **Model Weights**: Place LongCat model weights in a directory accessible to ComfyUI.
+
+## Usage in ComfyUI
+- Put this plugin folder into ComfyUI `custom_nodes`, then restart.
+- Nodes appear under category `longcat`.
+- **Workflow**:
+    1. Load model with `LongCatCheckpointLoader`.
+    2. Connect `CLIP` to `TextEncodeLongCatImage` (T2I) or `TextEncodeLongCatImageEdit` (Edit).
+    3. Connect `VAE` to `VAEEncodeLongCat` (if using image input) and `VAEDecodeLongCat`.
+    4. Connect `MODEL`, `CONDITIONING` (Positive/Negative), and `LATENT` to `LongCatSampler`.
+- Set `model_path` to your downloaded LongCat checkpoint; on CUDA you may enable `cpu_offload` to save VRAM.
+
+## Development
+- Tests: `pytest tests/test_nodes.py`
+- Key modules: `nodes.py`, `longcat_image/pipelines/*`, `longcat_image/models/longcat_image_dit.py`.
+
+## To-Do
+- [ ] Add robust fallback when prompt rewriting model lacks `generate`.
+- [ ] Guard `LongCatImageEditPipeline` for `image=None` misuse.
+- [ ] Expand docs with example ComfyUI workflows and screenshots.
+- [ ] Provide sample model download script and config guidance.
+- [ ] Add more automated tests (pipeline smoke tests, dtype/device matrix).
+
+## Credits
+- LongCat base model & pipelines: original LongCat project (LongCat team).
+- **diffusers** (Hugging Face) for pipeline scaffolding.
+- **transformers** (Hugging Face) for text and vision encoders.
+- **accelerate** for optional CPU/GPU offload helpers.
+- **PyTorch** for core tensor and model runtime.
+- **Pillow (PIL)** and **NumPy** for image/tensor conversions.
+- OpenAI / DeepSeek API usage scaffold in `misc/prompt_rewrite_api.py`.

README_CN.md

@@ -0,0 +1,53 @@
+# ComfyUI LongCat 插件
+
+[English](README.md)
+
+## 插件简介
+基于 Hugging Face diffusers 的 LongCat 文生图与图像编辑管线，对应的 ComfyUI 节点集。支持提示词改写、按比例自动计算分辨率、潜变量解码，方便在 ComfyUI 中直接使用。
+
+## 功能特性
+- **LongCatCheckpointLoader**：加载 LongCat 模型组件（Model, CLIP, VAE）。
+- **TextEncodeLongCatImage / TextEncodeLongCatImageEdit**：将文本提示词和参考图（编辑模式）编码为条件。
+- **VAEEncodeLongCat / VAEDecodeLongCat**：使用 LongCat VAE 将图像编码为潜变量，或将潜变量解码为图像。
+- **LongCatSampler**：使用 LongCat Transformer 进行采样（去噪）。
+- **LongCatImageSizeScale**：按目标像素面积缩放，并将宽高对齐到 16 的倍数。
+
+## 安装
+1.  **复制文件夹**：将 `comfyui_longcat` 文件夹复制到 ComfyUI 的 `custom_nodes` 目录下。
+    *   例如：`ComfyUI/custom_nodes/comfyui_longcat`
+2.  **安装依赖**：确保 ComfyUI 环境中安装了必要的包：
+    ```bash
+    pip install -r requirements.txt
+    ```
+    （注：`requirements.txt` 位于本仓库根目录，或查看 `setup.py` 中的列表）。
+3.  **模型权重**：将 LongCat 模型权重放入 ComfyUI 可访问的目录。
+
+## 在 ComfyUI 中使用
+- 将本插件目录放入 ComfyUI 的 `custom_nodes`，重启后生效。
+- 节点位于分类 `longcat`。
+- **工作流**：
+    1. 使用 `LongCatCheckpointLoader` 加载模型。
+    2. 连接 `CLIP` 到 `TextEncodeLongCatImage` (文生图) 或 `TextEncodeLongCatImageEdit` (编辑)。
+    3. 连接 `VAE` 到 `VAEEncodeLongCat` (如有图像输入) 和 `VAEDecodeLongCat`。
+    4. 连接 `MODEL`、`CONDITIONING` (正向/负向) 和 `LATENT` 到 `LongCatSampler`。
+- `model_path` 指向已下载的 LongCat 权重；在 CUDA 上可开启 `cpu_offload` 以节省显存。
+
+## 开发
+- 测试：`pytest tests/test_nodes.py`
+- 关键模块：`nodes.py`、`longcat_image/pipelines/*`、`longcat_image/models/longcat_image_dit.py`。
+
+## 待办
+- [ ] 当提示词改写模型缺少 `generate` 时增加可靠回退。
+- [ ] 为 `LongCatImageEditPipeline` 增加 `image=None` 误用保护。
+- [ ] 扩展文档，提供 ComfyUI 流程示例与截图。
+- [ ] 提供模型下载脚本与配置指引。
+- [ ] 增加更多自动化测试（管线冒烟、dtype/device 组合）。
+
+## 致谢
+- LongCat 基础模型与管线：原 LongCat 项目（LongCat 团队）。
+- **diffusers**（Hugging Face）提供管线框架。
+- **transformers**（Hugging Face）用于文本与视觉编码器。
+- **accelerate** 提供可选的 CPU/GPU offload。
+- **PyTorch** 提供核心张量与运行时。
+- **Pillow (PIL)** 与 **NumPy** 用于图像/张量转换。
+- `misc/prompt_rewrite_api.py` 中的 OpenAI/DeepSeek API 使用脚手架。

__init__.py

@@ -0,0 +1,31 @@
+from .nodes import (
+    LongCatCheckpointLoader,
+    TextEncodeLongCatImage,
+    TextEncodeLongCatImageEdit,
+    VAEEncodeLongCat,
+    VAEDecodeLongCat,
+    LongCatSampler,
+    LongCatImageSizeScale,
+)
+
+NODE_CLASS_MAPPINGS = {
+    "LongCatCheckpointLoader": LongCatCheckpointLoader,
+    "TextEncodeLongCatImage": TextEncodeLongCatImage,
+    "TextEncodeLongCatImageEdit": TextEncodeLongCatImageEdit,
+    "VAEEncodeLongCat": VAEEncodeLongCat,
+    "VAEDecodeLongCat": VAEDecodeLongCat,
+    "LongCatSampler": LongCatSampler,
+    "LongCatImageSizeScale": LongCatImageSizeScale,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LongCatCheckpointLoader": "LongCat Checkpoint Loader",
+    "TextEncodeLongCatImage": "LongCat Text Encode (T2I)",
+    "TextEncodeLongCatImageEdit": "LongCat Text Encode (Edit)",
+    "VAEEncodeLongCat": "LongCat VAE Encode",
+    "VAEDecodeLongCat": "LongCat VAE Decode",
+    "LongCatSampler": "LongCat Sampler",
+    "LongCatImageSizeScale": "LongCat Image Size Scale",
+}
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

longcat_image/dataset/__init__.py

@@ -0,0 +1,2 @@
+from .data_utils import MULTI_RESOLUTION_MAP
+from .sampler import MultiResolutionDistributedSampler

longcat_image/dataset/data_utils.py

@@ -0,0 +1,34 @@
+MULTI_ASPECT_RATIO_1024 = {
+    '0.5': [704., 1408.], '0.52': [704., 1344.],
+    '0.57': [768., 1344.], '0.6': [768., 1280.], '0.68': [832., 1216.], '0.72': [832., 1152.],
+    '0.78': [896., 1152.], '0.82': [896., 1088.], '0.88': [960., 1088.], '0.94': [960., 1024.],
+    '1.0':  [1024., 1024.], '1.07': [1024.,  960.], '1.13': [1088.,  960.], '1.21': [1088.,  896.],
+    '1.29': [1152.,  896.], '1.38': [1152.,  832.], '1.46': [1216.,  832.], '1.67': [1280.,  768.],
+    '1.75': [1344.,  768.], '2.0':  [1408.,  704.]
+}
+
+
+MULTI_ASPECT_RATIO_512 = {
+    '0.5': [352.0, 704.0], '0.52': [352.0, 672.0],
+    '0.57': [384.0, 672.0], '0.6': [384.0, 640.0], '0.68': [416.0, 608.0], '0.72': [416.0, 576.0],
+    '0.78': [448.0, 576.0], '0.82': [448.0, 544.0], '0.88': [480.0, 544.0], '0.94': [480.0, 512.0],
+    '1.0': [512.0, 512.0], '1.07': [512.0, 480.0], '1.13': [544.0, 480.0], '1.21': [544.0, 448.0],
+    '1.29': [576.0, 448.0], '1.38': [576.0, 416.0], '1.46': [608.0, 416.0], '1.67': [640.0, 384.0],
+    '1.75': [672.0, 384.0], '2.0': [704.0, 352.0]
+}
+
+
+MULTI_ASPECT_RATIO_256 = {
+    '0.5': [176.0, 352.0], '0.52': [176.0, 336.0],
+    '0.57': [192.0, 336.0], '0.6': [192.0, 320.0], '0.68': [208.0, 304.0], '0.72': [208.0, 288.0],
+    '0.78': [224.0, 288.0], '0.82': [224.0, 272.0], '0.88': [240.0, 272.0], '0.94': [240.0, 256.0],
+    '1.0': [256.0, 256.0], '1.07': [256.0, 240.0], '1.13': [272.0, 240.0], '1.21': [272.0, 224.0],
+    '1.29': [288.0, 224.0], '1.38': [288.0, 208.0], '1.46': [304.0, 208.0], '1.67': [320.0, 192.0],
+    '1.75': [336.0, 192.0], '2.0': [352.0, 176.0]
+}
+
+MULTI_RESOLUTION_MAP = {
+    'mar_256': MULTI_ASPECT_RATIO_256,
+    'mar_512': MULTI_ASPECT_RATIO_512,
+    'mar_1024': MULTI_ASPECT_RATIO_1024,
+}

longcat_image/dataset/sampler.py

@@ -0,0 +1,111 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+import torch.distributed as dist
+import numpy as np
+import copy
+from torch.utils.data import IterableDataset
+
+from longcat_image.utils.dist_utils import get_world_size, get_rank, get_local_rank
+
+class MultiResolutionDistributedSampler(torch.utils.data.Sampler):
+    def __init__(self,
+                 batch_size: int,
+                 dataset: IterableDataset,
+                 data_resolution_infos: List,
+                 bucket_info: dict,
+                 num_replicas: int = None,
+                 rank: int = None,
+                 seed: int = 888,
+                 epoch: int = 0,
+                 shuffle: bool = True):
+
+        if not dist.is_available():
+            num_replicas = 1
+            rank = 0
+        else:
+            num_replicas = get_world_size()
+            rank = get_rank()
+
+        self.len_items = len(dataset)
+        bucket_info = {float(b): bucket_info[b] for b in bucket_info.keys()}
+        self.aspect_ratios = np.array(sorted(list(bucket_info.keys())))
+        self.resolutions = np.array([bucket_info[aspect] for aspect in self.aspect_ratios])
+
+        self.batch_size = batch_size
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = epoch
+        self.shuffle = shuffle
+        self.seed = seed
+        self.cur_rank_index = []
+        self.rng = np.random.RandomState(seed+self.epoch)
+        self.global_batch_size = batch_size*num_replicas
+        self.data_resolution_infos = np.array(data_resolution_infos, dtype=np.float32)
+        print(f'num_replicas {num_replicas}, cur rank {rank}!!!')
+
+        self.split_to_buckets()
+        self.num_samples = len(dataset)//num_replicas
+
+    def split_to_buckets(self):
+        self.buckets = {}
+        self._buckets_bak = {}
+        data_aspect_ratio = self.data_resolution_infos[:,0]*1.0/self.data_resolution_infos[:, 1]
+        bucket_id = np.abs(data_aspect_ratio[:, None] - self.aspect_ratios).argmin(axis=1)
+        for i in range(len(self.aspect_ratios)):
+            self.buckets[i] = np.where(bucket_id == i)[0]
+            self._buckets_bak[i] = np.where(bucket_id == i)[0]
+        for k, v in self.buckets.items():
+            print(f'bucket {k}, resolutions {self.resolutions[k]}, sampler nums {len(v)}!!!')
+
+    def get_batch_index(self):
+        success_flag = False
+        while not success_flag:
+            bucket_ids = list(self.buckets.keys())
+            bucket_probs = [len(self.buckets[bucket_id]) for bucket_id in bucket_ids]
+            bucket_probs = np.array(bucket_probs, dtype=np.float32)
+            bucket_probs = bucket_probs / bucket_probs.sum()
+            bucket_ids = np.array(bucket_ids, dtype=np.int64)
+            chosen_id = int(self.rng.choice(bucket_ids, 1, p=bucket_probs)[0])
+            if len(self.buckets[chosen_id]) < self.global_batch_size:
+                del self.buckets[chosen_id]
+                continue
+            batch_data = self.buckets[chosen_id][:self.global_batch_size]
+            batch_data = (batch_data, self.resolutions[chosen_id])
+            self.buckets[chosen_id] = self.buckets[chosen_id][self.global_batch_size:]
+            if len(self.buckets[chosen_id]) == 0:
+                del self.buckets[chosen_id]
+            success_flag = True
+            assert bool(self.buckets), 'There is not enough data in the current epoch.'
+        return batch_data
+
+    def shuffle_bucker_index(self):
+        self.rng = np.random.RandomState(self.seed+self.epoch)
+        self.buckets = copy.deepcopy(self._buckets_bak)
+        for bucket_id in self.buckets.keys():
+            self.rng.shuffle(self.buckets[bucket_id])
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            if len(self.cur_rank_index) == 0:
+                global_batch_index, target_resolutions = self.get_batch_index()
+                self.cur_rank_index = list(map(
+                    int, global_batch_index[self.batch_size*self.rank:self.batch_size*(self.rank+1)]))
+                self.resolution = list(map(int, target_resolutions))
+            data_index = self.cur_rank_index.pop(0)
+            return (data_index, self.resolution)
+
+        except Exception as e:
+            self.epoch += 1
+            self.shuffle_bucker_index()
+            print(f'get error {e}.')
+            raise StopIteration
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch

longcat_image/models/__init__.py

@@ -0,0 +1 @@
+from .longcat_image_dit import LongCatImageTransformer2DModel

longcat_image/models/longcat_image_dit.py

@@ -0,0 +1,231 @@
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models.transformers.transformer_flux import \
+    FluxTransformerBlock, FluxSingleTransformerBlock, \
+    AdaLayerNormContinuous, Transformer2DModelOutput
+from diffusers.models.embeddings import Timesteps, TimestepEmbedding,FluxPosEmbed
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from accelerate.logging import get_logger
+from diffusers.loaders import PeftAdapterMixin
+
+logger = get_logger(__name__, log_level="INFO")
+
+
+class TimestepEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timestep, hidden_dtype):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))  # (N, D)
+
+        return timesteps_emb
+    
+
+class LongCatImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin ):
+    """
+    The Transformer model introduced in Flux.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @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 = 3584,
+        pooled_projection_dim: int = 3584,
+        axes_dims_rope: List[int] = [16, 56, 56],
+    ):
+        super().__init__()
+        self.out_channels = in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.pooled_projection_dim = pooled_projection_dim
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+
+        self.time_embed = TimestepEmbeddings(embedding_dim=self.inner_dim)
+
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
+        self.x_embedder = torch.nn.Linear(in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for i in range(num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(
+            self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+        self.initialize_weights()
+
+        self.use_checkpoint = [True] * num_layers
+        self.use_single_checkpoint = [True] * num_single_layers
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        return_dict: bool = True,
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        """
+        The  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.
+            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.
+            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.
+        """
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+
+        temb = self.time_embed( timestep, hidden_states.dtype )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if txt_ids.ndim == 3:
+            logger.warning(
+                "Passing `txt_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            logger.warning(
+                "Passing `img_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            img_ids = img_ids[0]
+
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        image_rotary_emb = self.pos_embed(ids)
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing and self.use_checkpoint[index_block]:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing and self.use_single_checkpoint[index_block]:
+                encoder_hidden_states,hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.bias, 0)
+
+        # Initialize caption embedding MLP:
+        nn.init.normal_(self.context_embedder.weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in blocks:
+        for block in self.transformer_blocks:
+            nn.init.constant_(block.norm1.linear.weight, 0)
+            nn.init.constant_(block.norm1.linear.bias, 0)
+            nn.init.constant_(block.norm1_context.linear.weight, 0)
+            nn.init.constant_(block.norm1_context.linear.bias, 0)
+
+        for block in self.single_transformer_blocks:
+            nn.init.constant_(block.norm.linear.weight, 0)
+            nn.init.constant_(block.norm.linear.bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.norm_out.linear.weight, 0)
+        nn.init.constant_(self.norm_out.linear.bias, 0)
+        nn.init.constant_(self.proj_out.weight, 0)
+        nn.init.constant_(self.proj_out.bias, 0)
+

longcat_image/pipelines/__init__.py

@@ -0,0 +1,2 @@
+from .pipeline_longcat_image import LongCatImagePipeline
+from .pipeline_longcat_image_edit import LongCatImageEditPipeline

longcat_image/pipelines/pipeline_longcat_image.py

@@ -0,0 +1,576 @@
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Callable, Dict, List, Optional, Union
+import json
+import numpy as np
+import torch
+import re
+
+from transformers import (
+    CLIPImageProcessor,
+    CLIPVisionModelWithProjection,
+)
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+)
+from transformers import AutoTokenizer, AutoModel, AutoProcessor
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+from longcat_image.utils.model_utils import split_quotation, prepare_pos_ids, calculate_shift, retrieve_timesteps, optimized_scale
+from longcat_image.models.longcat_image_dit import LongCatImageTransformer2DModel
+from longcat_image.pipelines.pipeline_output import LongCatImagePipelineOutput
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+
+SYSTEM_PROMPT_EN = """
+You are a prompt engineering expert for text-to-image models. Since text-to-image models have limited capabilities in understanding user prompts, you need to identify the core theme and intent of the user's input and improve the model's understanding accuracy and generation quality through optimization and rewriting. The rewrite must strictly retain all information from the user's original prompt without deleting or distorting any details.
+Specific requirements are as follows:
+1. The rewrite must not affect any information expressed in the user's original prompt; the rewritten prompt should use coherent natural language, avoid low-information redundant descriptions, and keep the rewritten prompt length as concise as possible.
+2. Ensure consistency between input and output languages: Chinese input yields Chinese output, and English input yields English output. The rewritten token count should not exceed 512.
+3. The rewritten description should further refine subject characteristics and aesthetic techniques appearing in the original prompt, such as lighting and textures.
+4. If the original prompt does not specify an image style, ensure the rewritten prompt uses a **realistic photography style**. If the user specifies a style, retain the user's style.
+5. When the original prompt requires reasoning to clarify user intent, use logical reasoning based on world knowledge to convert vague abstract descriptions into specific tangible objects (e.g., convert "the tallest animal" to "a giraffe").
+6. When the original prompt requires text generation, please use double quotes to enclose the text part (e.g., `"50% OFF"`).
+7. When the original prompt requires generating text-heavy scenes like webpages, logos, UIs, or posters, and no specific text content is specified, you need to infer appropriate text content and enclose it in double quotes. For example, if the user inputs: "A tourism flyer with a grassland theme," it should be rewritten as: "A tourism flyer with the image title 'Grassland'."
+8. When negative words exist in the original prompt, ensure the rewritten prompt does not contain negative words. For example, "a lakeside without boats" should be rewritten such that the word "boat" does not appear at all.
+9. Except for text content explicitly requested by the user, **adding any extra text content is prohibited**.
+Here are examples of rewrites for different types of prompts:
+# Examples (Few-Shot Learning)
+  1. User Input: An animal with nine lives.
+    Rewrite Output: A cat bathed in soft sunlight, its fur soft and glossy. The background is a comfortable home environment with light from the window filtering through curtains, creating a warm light and shadow effect. The shot uses a medium distance perspective to highlight the cat's leisurely and stretched posture. Light cleverly hits the cat's face, emphasizing its spirited eyes and delicate whiskers, adding depth and affinity to the image.
+  2. User Input: Create an anime-style tourism flyer with a grassland theme.
+    Rewrite Output: In the lower right of the center, a short-haired girl sits sideways on a gray, irregularly shaped rock. She wears a white short-sleeved dress and brown flat shoes, holding a bunch of small white flowers in her left hand, smiling with her legs hanging naturally. The girl has dark brown shoulder-length hair with bangs covering her forehead, brown eyes, and a slightly open mouth. The rock surface has textures of varying depths. To the girl's left and front is lush grass, with long, yellow-green blades, some glowing golden in the sunlight. The grass extends into the distance, forming rolling green hills that fade in color as they recede. The sky occupies the upper half of the picture, pale blue dotted with a few fluffy white clouds. In the upper left corner, there is a line of text in italic, dark green font reading "Explore Nature's Peace". Colors are dominated by green, blue, and yellow, fluid lines, and distinct light and shadow contrast, creating a quiet and comfortable atmosphere.
+  3. User Input: A Christmas sale poster with a red background, promoting a Buy 1 Get 1 Free milk tea offer.
+    Rewrite Output: The poster features an overall red tone, embellished with white snowflake patterns on the top and left side. The upper right features a bunch of holly leaves with red berries and a pine cone. In the upper center, golden 3D text reads "Christmas Heartwarming Feedback" centered, along with red bold text "Buy 1 Get 1". Below, two transparent cups filled with bubble tea are placed side by side; the tea is light brown with dark brown pearls scattered at the bottom and middle. Below the cups, white snow piles up, decorated with pine branches, red berries, and pine cones. A blurry Christmas tree is faintly visible in the lower right corner. The image has high clarity, accurate text content, a unified design style, a prominent Christmas theme, and a reasonable layout, providing strong visual appeal.
+  4. User Input: A woman indoors shot in natural light, smiling with arms crossed, showing a relaxed and confident posture.
+    Rewrite Output: The image features a young Asian woman with long dark brown hair naturally falling over her shoulders, with some strands illuminated by light, showing a soft sheen. Her features are delicate, with long eyebrows, bright and spirited dark brown eyes looking directly at the camera, revealing peace and confidence. She has a high nose bridge, full lips with nude lipstick, and corners of the mouth slightly raised in a faint smile. Her skin is fair, with cheeks and collarbones illuminated by warm light, showing a healthy ruddiness. She wears a black spaghetti strap tank top revealing graceful collarbone lines, and a thin gold necklace with small beads and metal bars glinting in the light. Her outer layer is a beige knitted cardigan, soft in texture with visible knitting patterns on the sleeves. Her arms are crossed over her chest, hands covered by the cardigan sleeves, in a relaxed posture. The background is a pure dark brown without extra decoration, making the figure the absolute focus. The figure is located in the center of the frame. Light enters from the upper right, creating bright spots on her left cheek, neck, and collarbone, while the right side is slightly shadowed, creating a three-dimensional and soft tone. Image details are clear, showcasing skin texture, hair, and clothing materials well. Colors are dominated by warm tones, with the combination of beige and dark brown creating a warm and comfortable atmosphere. The overall style is natural, elegant, and artistic.
+  5. User Input: Create a series of images showing the growth process of an apple from seed to fruit. The series should include four stages: 1. Sowing, 2. Seedling growth, 3. Plant maturity, 4. Fruit harvesting.
+    Rewrite Output: A 4-panel exquisite illustration depicting the growth process of an apple, capturing each stage precisely and clearly. 1. "Sowing": A close-up shot of a hand gently placing a small apple seed into fertile dark soil, with visible soil texture and the seed's smooth surface. The background is a soft-focus garden dotted with green leaves and sunlight filtering through. 2. "Seedling Growth": A young apple sapling breaks through the soil, stretching tender green leaves toward the sky. The scene is set in a vibrant garden illuminated by warm golden light, highlighting the seedling's delicate structure. 3. "Plant Maturity": A mature apple tree, lush with branches and leaves, covered in tender green foliage and developing small apples. The background is a vibrant orchard under a clear blue sky, with dappled sunlight creating a peaceful atmosphere. 4. "Fruit Harvesting": A hand reaches into the tree to pick a ripe red apple, its smooth skin glistening in the sun. The scene shows the abundance of the orchard, with baskets of apples in the background, giving a sense of fulfillment. Each illustration uses a realistic style, focusing on details and harmonious colors to showcase the natural beauty and development of the apple's life cycle.
+  6. User Input: If 1 represents red, 2 represents green, 3 represents purple, and 4 represents yellow, please generate a four-color rainbow based on this rule. The color order from top to bottom is 3142.
+    Rewrite Output: The image consists of four horizontally arranged colored stripes, ordered from top to bottom as purple, red, yellow, and green. A white number is centered on each stripe. The top purple stripe features the number "3", the red stripe below it has the number "1", the yellow stripe further down has the number "4", and the bottom green stripe has the number "2". All numbers use a sans-serif font in pure white, forming a sharp contrast with the background colors to ensure good readability. The stripes have high color saturation and a slight texture. The overall layout is simple and clear, with distinct visual effects and no extra decorative elements, emphasizing the numerical information. The image is high definition, with accurate colors and a consistent style, offering strong visual appeal.
+  7. User Input: A stone tablet carved with "Guan Guan Ju Jiu, On the River Isle", natural light, background is a Chinese garden.
+    Rewrite Output: An ancient stone tablet carved with "Guan Guan Ju Jiu, On the River Isle", the surface covered with traces of time, the writing clear and deep. Natural light falls from above, softly illuminating every detail of the stone tablet and enhancing its sense of history. The background is an elegant Chinese garden featuring lush bamboo forests, winding paths, and quiet pools, creating a serene and distant atmosphere. The overall picture uses a realistic style with rich details and natural light and shadow effects, highlighting the cultural heritage of the stone tablet and the classical beauty of the garden.
+# Output Format
+Please directly output the rewritten and optimized Prompt content. Do not include any explanatory language or JSON formatting, and do not add opening or closing quotes yourself."""
+
+
+SYSTEM_PROMPT_ZH = """
+你是一名文生图模型的prompt engineering专家。由于文生图模型对用户prompt的理解能力有限，你需要识别用户输入的核心主题和意图，并通过优化改写提升模型的理解准确性和生成质量。改写必须严格保留用户原始prompt的所有信息，不得删减或曲解任何细节。
+具体要求如下：
+1. 改写不能影响用户原始prompt里表达的任何信息，改写后的prompt应该使用连贯的自然语言表达,不要出现低信息量的冗余描述，尽可能保持改写后prompt长度精简。
+2. 请确保输入和输出的语言类型一致，中文输入中文输出，英文输入英文输出，改写后的token数量不要超过512个;
+3. 改写后的描述应当进一步完善原始prompt中出现的主体特征、美学技巧，如打光、纹理等；
+4. 如果原始prompt没有指定图片风格时，确保改写后的prompt使用真实摄影风格，如果用户指定了图片风格，则保留用户风格；
+5. 当原始prompt需要推理才能明确用户意图时，根据世界知识进行适当逻辑推理，将模糊抽象描述转化为具体指向事物（例：将"最高的动物"转化为"一头长颈鹿"）。
+6. 当原始prompt需要生成文字时，请使用双引号圈定文字部分，例：`"限时5折"`）。
+7. 当原始prompt需要生成网页、logo、ui、海报等文字场景时，且没有指定具体的文字内容时，需要推断出合适的文字内容，并使用双引号圈定，如用户输入：一个旅游宣传单，以草原为主题。应该改写成：一个旅游宣传单，图片标题为“草原”。
+8. 当原始prompt中存在否定词时，需要确保改写后的prompt不存在否定词，如没有船的湖边，改写后的prompt不能出现船这个词汇。
+9. 除非用户指定生成品牌logo，否则不要增加额外的品牌logo.
+10. 除了用户明确要求书写的文字内容外，**禁止增加任何额外的文字内容**。
+以下是针对不同类型prompt改写的示例：
+
+# Examples (Few-Shot Learning)
+  1. 用户输入: 九条命的动物。
+    改写输出: 一只猫，被柔和的阳光笼罩着，毛发柔软而富有光泽。背景是一个舒适的家居环境，窗外的光线透过窗帘，形成温馨的光影效果。镜头采用中距离视角，突出猫悠闲舒展的姿态。光线巧妙地打在猫的脸部，强调它灵动的眼睛和精致的胡须，增加画面的层次感与亲和力。
+  2. 用户输入: 制作一个动画风格的旅游宣传单，以草原为主题。
+    改写输出: 画面中央偏右下角，一个短发女孩侧身坐在灰色的不规则形状岩石上，她穿着白色短袖连衣裙和棕色平底鞋，左手拿着一束白色小花，面带微笑，双腿自然垂下。女孩的头发为深棕色，齐肩短发，刘海覆盖额头，眼睛呈棕色，嘴巴微张。岩石表面有深浅不一的纹理。女孩的左侧和前方是茂盛的草地，草叶细长，呈黄绿色，部分草叶在阳光下泛着金色的光芒，仿佛被阳光照亮。草地向远处延伸，形成连绵起伏的绿色山丘，山丘的颜色由近及远逐渐变浅。天空占据了画面的上半部分，呈淡蓝色，点缀着几朵白色蓬松的云彩。画面的左上角有一行文字，文字内容是斜体、深绿色的“Explore Nature's Peace”。色彩以绿色、蓝色和黄色为主，线条流畅，光影明暗对比明显，营造出一种宁静、舒适的氛围。
+  3. 用户输入: 一张以红色为背景的圣诞节促销海报，主要宣传奶茶买一送一的优惠活动。
+    改写输出: 海报整体呈现红色调，上方和左侧点缀着白色雪花图案，右上方有一束冬青叶和红色浆果，以及一个松果。海报中央偏上位置，金色立体字样“圣诞节 暖心回馈”居中排列，和红色粗体字“买1送1”。海报下方，两个装满珍珠奶茶的透明杯子并排摆放，杯中奶茶呈浅棕色，底部和中间散布着深棕色珍珠。杯子下方，堆积着白色雪花，雪花上装饰着松枝、红色浆果和松果。右下角隐约可见一棵模糊的圣诞树。图片清晰度高，文字内容准确，整体设计风格统一，圣诞主题突出，排版布局合理，具有较强的视觉吸引力。
+  4. 用户输入: 一位女性在室内以自然光线拍摄，她面带微笑，双臂交叉，展现出轻松自信的姿态。
+    改写输出: 画面中是一位年轻的亚洲女性，她拥有深棕色的长发，发丝自然地垂落在双肩，部分发丝被光线照亮，呈现出柔和的光泽。她的五官精致，眉毛修长，眼睛明亮有神，瞳孔呈深棕色，眼神直视镜头，流露出平和与自信。鼻梁挺拔，嘴唇丰满，涂有裸色系唇膏，嘴角微微上扬，展现出浅浅的微笑。她的肤色白皙，脸颊和锁骨处被暖色调的光线照亮，呈现出健康的红润感。她穿着一件黑色的细吊带背心，肩带纤细，露出优美的锁骨线条。脖颈上佩戴着一条金色的细项链，项链由小珠子和几个细长的金属条组成，在光线下闪烁着光泽。她的外搭是一件米黄色的针织开衫，材质柔软，袖子部分有明显的针织纹理。她双臂交叉在胸前，双手被开衫的袖子覆盖，姿态放松。背景是纯粹的深棕色，没有多余的装饰，使得人物成为画面的绝对焦点。人物位于画面中央。光线从画面的右上方射入，在人物的左侧脸颊、脖颈和锁骨处形成明亮的光斑，右侧则略显阴影，营造出立体感和柔和的影调。图像细节清晰，人物的皮肤纹理、发丝以及衣物材质都得到了很好的展现。色彩以暖色调为主，米黄色和深棕色的搭配营造出温馨舒适的氛围。整体呈现出一种自然、优雅且富有亲和力的艺术风格。
+  5. 用户输入：创作一系列图片，展现苹果从种子到结果的生长过程。该系列图片应包含以下四个阶段：1. 播种，2. 幼苗生长，3. 植物成熟，4. 果实采摘。
+    改写输出：一个4宫格的精美插图，描绘苹果的生长过程，精确清晰地捕捉每个阶段。1.“播种”：特写镜头，一只手轻轻地将一颗小小的苹果种子放入肥沃的深色土壤中，土壤的纹理和种子光滑的表面清晰可见。背景是花园的柔焦画面，点缀着绿色的树叶和透过树叶洒下的阳光。2.“幼苗生长”：一棵幼小的苹果树苗破土而出，嫩绿的叶子向天空舒展。场景设定在一个生机勃勃的花园中，温暖的金光照亮了它。幼苗的纤细结构。3.“植物的成熟”：一棵成熟的苹果树，枝繁叶茂，挂满了嫩绿的叶子和正在萌发的小苹果。背景是一片生机勃勃的果园，湛蓝的天空下，斑驳的阳光营造出宁静祥和的氛围。4.“采摘果实”：一只手伸向树上，摘下一个成熟的红苹果，苹果光滑的果皮在阳光下闪闪发光。画面展现了果园的丰收景象，背景中摆放着一篮篮的苹果，给人一种圆满满足的感觉。每幅插图都采用写实风格，注重细节，色彩和谐，展现了苹果生命周期的自然之美和发展过程。
+  6. 用户输入： 如果1代表红色，2代表绿色，3代表紫色，4代表黄色，请按照此规则生成四色彩虹。它的颜色顺序从上到下是3142
+    改写输出：图片由四个水平排列的彩色条纹组成，从上到下依次为紫色、红色、黄色和绿色。每个条纹上都居中放置一个白色数字。最上方的紫色条纹上是数字“3”，其下方红色条纹上是数字“1”，再下方黄色条纹上是数字“4”，最下方的绿色条纹上是数字“2”。所有数字均采用无衬线字体，颜色为纯白色，与背景色形成鲜明对比，确保了良好的可读性。条纹的颜色饱和度高，且带有轻微的纹理感，整体排版简洁明了，视觉效果清晰，没有多余的装饰元素，强调了数字信息本身。图片整体清晰度高，色彩准确，风格一致，具有较强的视觉吸引力。
+  7. 用户输入：石碑上刻着“关关雎鸠，在河之洲”，自然光照，背景是中式园林
+    改写输出：一块古老的石碑上刻着“关关雎鸠，在河之洲”，石碑表面布满岁月的痕迹，字迹清晰而深刻。自然光线从上方洒下，柔和地照亮石碑的每一个细节，增强了其历史感。背景是一座典雅的中式园林，园林中有翠绿的竹林、蜿蜒的小径和静谧的水池，营造出一种宁静而悠远的氛围。整体画面采用写实风格，细节丰富，光影效果自然，突出了石碑的文化底蕴和园林的古典美。
+# 输出格式
+请直接输出改写优化后的 Prompt 内容，不要包含任何解释性语言或 JSON 格式，不要自行添加开头或结尾的引号。
+"""
+
+def get_prompt_language(prompt):
+    pattern = re.compile(r'[\u4e00-\u9fff]')
+    if bool(pattern.search(prompt)):
+      return 'zh'
+    return 'en'
+
+
+class LongCatImagePipeline(
+    DiffusionPipeline,
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+):
+    r"""
+    The pipeline for text-to-image generation.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+    """
+
+    model_cpu_offload_seq = "text_encoder->image_encoder->transformer->vae"
+    _optional_components = ["image_encoder", "feature_extractor", "text_processor"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: AutoModel,
+        tokenizer: AutoTokenizer,
+        text_processor: AutoProcessor,
+        transformer: LongCatImageTransformer2DModel,
+        image_encoder: CLIPVisionModelWithProjection,
+        feature_extractor: CLIPImageProcessor = None,
+        
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+            text_processor=text_processor,
+        )
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        # Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        
+        self.prompt_template_encode_prefix = '<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n'
+        self.prompt_template_encode_suffix = '<|im_end|>\n<|im_start|>assistant\n'
+        self.prompt_template_encode_start_idx = 36
+        self.prompt_template_encode_end_idx = 5
+        self.default_sample_size = 128
+        self.max_tokenizer_len = 512
+
+
+    @torch.inference_mode()
+    def rewire_prompt(self, prompt, device):
+        language = get_prompt_language(prompt)
+        if language == 'zh':
+            question = SYSTEM_PROMPT_ZH + f"\n用户输入为：{prompt}\n改写后的prompt为："
+        else:
+            question = SYSTEM_PROMPT_EN + f"\nUser Input: {prompt}\nRewritten prompt:"
+            
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": question},
+                ],
+            }
+        ]
+        # Preparation for inference
+        text = self.text_processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        inputs = self.text_processor(
+            text=[text],padding=True,return_tensors="pt",)
+        inputs = inputs.to(device)
+
+        generated_ids = self.text_encoder.generate(**inputs, max_new_tokens=self.max_tokenizer_len)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = self.text_processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        rewrite_prompt= output_text
+        return rewrite_prompt
+    
+    @torch.inference_mode()
+    def encode_prompt(self,
+                      prompts,
+                      device,
+                      dtype):
+
+        prompts = [prompt.strip('"') if prompt.startswith('"') and prompt.endswith('"') else prompt for prompt in prompts]
+        all_tokens = []
+        for clean_prompt_sub, matched in split_quotation(prompts[0]):
+            if matched:
+                for sub_word in clean_prompt_sub:
+                    tokens = self.tokenizer(sub_word, add_special_tokens=False)['input_ids']
+                    all_tokens.extend(tokens)
+            else:
+                tokens = self.tokenizer(clean_prompt_sub, add_special_tokens=False)['input_ids']
+                all_tokens.extend(tokens)
+
+        all_tokens = all_tokens[:self.max_tokenizer_len]
+        text_tokens_and_mask = self.tokenizer.pad(
+            {'input_ids': [all_tokens]},
+            max_length=self.max_tokenizer_len,
+            padding='max_length',
+            return_attention_mask=True,
+            return_tensors='pt')
+
+        prefix_tokens = self.tokenizer(self.prompt_template_encode_prefix, add_special_tokens=False)['input_ids']
+        suffix_tokens = self.tokenizer(self.prompt_template_encode_suffix, add_special_tokens=False)['input_ids']
+        prefix_tokens_mask = torch.tensor( [1]*len(prefix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype )
+        suffix_tokens_mask = torch.tensor( [1]*len(suffix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype )
+
+        prefix_tokens = torch.tensor(prefix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+        suffix_tokens = torch.tensor(suffix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+                
+        input_ids = torch.cat( (prefix_tokens, text_tokens_and_mask.input_ids[0], suffix_tokens), dim=-1 )
+        attention_mask = torch.cat( (prefix_tokens_mask, text_tokens_and_mask.attention_mask[0], suffix_tokens_mask), dim=-1 )
+
+        input_ids = input_ids.unsqueeze(0).to(self.device)
+        attention_mask = attention_mask.unsqueeze(0).to(self.device)
+
+        text_output = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True
+        )
+        # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size]
+        # clone to have a contiguous tensor
+        prompt_embeds = text_output.hidden_states[-1].detach()
+        prompt_embeds = prompt_embeds[:,self.prompt_template_encode_start_idx: -self.prompt_template_encode_end_idx ,:]
+
+        text_ids = prepare_pos_ids(modality_id=0,
+                                   type='text',
+                                   start=(0, 0),
+                                   num_token=prompt_embeds.shape[1]).to(device, dtype=dtype)
+
+        return prompt_embeds,  text_ids
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+        return latents
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+        latent_image_ids = prepare_pos_ids(modality_id=1,
+                                               type='image',
+                                               start=(self.max_tokenizer_len,
+                                                      self.max_tokenizer_len),
+                                               height=height//2,
+                                               width=width//2).to(device, dtype=torch.float64)
+        
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype), latent_image_ids
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator,device=device)
+        latents = latents.to(dtype=dtype)
+        latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+
+        return latents, latent_image_ids
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator,
+                                  List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        enable_cfg_renorm: Optional[bool] = True,
+        cfg_renorm_min: Optional[float] = 0.0,
+        enable_prompt_rewrite: Optional[bool] = True,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            enable_cfg_renorm: Whether to enable cfg_renorm. Enabling cfg_renorm will improve image quality,
+                but it may lead to a decrease in the stability of some image outputs..
+            cfg_renorm_min: The minimum value of the cfg_renorm_scale range (0-1).
+                cfg_renorm_min = 1.0, renorm has no effect, while cfg_renorm_min=0.0, the renorm range is larger.
+            enable_prompt_rewrite: whether to enable prompt rewrite.
+        Examples:
+
+        Returns:
+            [`~pipelines.LongCatImagePipelineOutput`] or `tuple`: [`~pipelines.LongCatImagePipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+            pixel_step= self.vae_scale_factor * 2
+            height = int( height/pixel_step )*pixel_step
+            width = int( width/pixel_step )*pixel_step
+        
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        
+        device = self._execution_device
+        if enable_prompt_rewrite:
+            prompt = self.rewire_prompt(prompt, device )
+        
+        negative_prompt = '' if negative_prompt is None else negative_prompt
+        negative_prompt = [negative_prompt]*num_images_per_prompt
+        prompt = [prompt]*num_images_per_prompt
+
+        prompt_embeds, text_ids = self.encode_prompt(prompt, device, dtype=torch.float64)
+        negative_prompt_embeds, negative_text_ids = self.encode_prompt(negative_prompt, device, dtype=torch.float64)
+
+        # 4. Prepare latent variables
+        num_channels_latents = 16
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1.0 / num_inference_steps,num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        guidance = None
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0).to(device)
+        else:
+            prompt_embeds = prompt_embeds.to(device)
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
+                
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    return_dict=False,
+                )[0]
+
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * \
+                            (noise_pred_text - noise_pred_uncond)
+                    
+                    if enable_cfg_renorm:
+                        cond_norm = torch.norm(noise_pred_text, dim=-1, keepdim=True)
+                        noise_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+                        scale = (cond_norm / (noise_norm + 1e-8)).clamp(min= cfg_renorm_min , max=1.0)
+                        noise_pred = noise_pred * scale
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            if latents.dtype != self.vae.dtype:
+                latents = latents.to(dtype=self.vae.dtype)
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        # self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return LongCatImagePipelineOutput(images=image)

longcat_image/pipelines/pipeline_longcat_image_edit.py

@@ -0,0 +1,512 @@
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Callable, Dict, List, Optional, Union
+import json
+import numpy as np
+import torch
+import math
+from transformers import (
+    CLIPImageProcessor,
+    CLIPVisionModelWithProjection,
+)
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+)
+from transformers import AutoTokenizer, AutoModel, AutoProcessor
+
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+from longcat_image.utils.model_utils import split_quotation, prepare_pos_ids, calculate_shift, retrieve_timesteps, optimized_scale
+from longcat_image.models.longcat_image_dit import LongCatImageTransformer2DModel
+from longcat_image.pipelines.pipeline_output import LongCatImagePipelineOutput
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def calculate_dimensions(target_area, ratio):
+    width = math.sqrt(target_area * ratio)
+    height = width / ratio
+
+    width = width if width % 16 == 0 else (width // 16 + 1) * 16
+    height = height if height % 16 == 0 else (height // 16 + 1) * 16
+
+    width = int(width)
+    height = int(height)
+
+    return width, height
+
+
+class LongCatImageEditPipeline(
+    DiffusionPipeline,
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+):
+    r"""
+    The pipeline for text-to-image generation.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+    """
+
+    model_cpu_offload_seq = "text_encoder->image_encoder->transformer->vae"
+    _optional_components = ["image_encoder", "feature_extractor", "text_processor"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: AutoModel,
+        tokenizer: AutoTokenizer,
+        text_processor: AutoProcessor,
+        transformer: LongCatImageTransformer2DModel,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+            text_processor=text_processor,
+        )
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        # Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.image_processor_vl = text_processor.image_processor
+        
+        self.image_token = "<|image_pad|>"
+        self.prompt_template_encode_prefix = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        self.prompt_template_encode_suffix = '<|im_end|>\n<|im_start|>assistant\n'
+        self.prompt_template_encode_start_idx = 67
+        self.prompt_template_encode_end_idx = 5
+        self.default_sample_size = 128
+        self.max_tokenizer_len = 512
+        self.latent_channels = 16
+
+
+    @torch.inference_mode()
+    def encode_prompt(self,
+                      image,
+                      prompts,
+                      device,
+                      dtype):
+        raw_vl_input = self.image_processor_vl(images=image,return_tensors="pt")
+        pixel_values = raw_vl_input['pixel_values']
+        image_grid_thw = raw_vl_input['image_grid_thw']
+
+        prompts = [prompt.strip('"') if prompt.startswith('"') and prompt.endswith('"') else prompt for prompt in prompts]
+        all_tokens = []
+        
+        for clean_prompt_sub, matched in split_quotation(prompts[0]):
+            if matched:
+                for sub_word in clean_prompt_sub:
+                    tokens = self.tokenizer(sub_word, add_special_tokens=False)['input_ids']
+                    all_tokens.extend(tokens)
+            else:
+                tokens = self.tokenizer(clean_prompt_sub, add_special_tokens=False)['input_ids']
+                all_tokens.extend(tokens)
+
+        all_tokens = all_tokens[:self.max_tokenizer_len]
+        text_tokens_and_mask = self.tokenizer.pad(
+            {'input_ids': [all_tokens]},
+            max_length=self.max_tokenizer_len,
+            padding='max_length',
+            return_attention_mask=True,
+            return_tensors='pt')
+        text = self.prompt_template_encode_prefix
+        
+        merge_length = self.image_processor_vl.merge_size**2
+        while self.image_token in text:
+            num_image_tokens = image_grid_thw.prod() // merge_length
+            text = text.replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1)
+        text = text.replace("<|placeholder|>", self.image_token)
+
+        prefix_tokens = self.tokenizer(text, add_special_tokens=False)['input_ids']
+        suffix_tokens = self.tokenizer(self.prompt_template_encode_suffix, add_special_tokens=False)['input_ids']
+        prefix_tokens_mask = torch.tensor([1]*len(prefix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype)
+        suffix_tokens_mask = torch.tensor([1]*len(suffix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype)
+
+        prefix_tokens = torch.tensor(prefix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+        suffix_tokens = torch.tensor(suffix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+        
+        input_ids = torch.cat((prefix_tokens, text_tokens_and_mask.input_ids[0], suffix_tokens), dim=-1 )
+        attention_mask = torch.cat((prefix_tokens_mask, text_tokens_and_mask.attention_mask[0], suffix_tokens_mask), dim=-1)
+
+        pixel_values = pixel_values.to(self.device)
+        image_grid_thw = image_grid_thw.to(self.device)
+
+        input_ids = input_ids.unsqueeze(0).to(self.device)
+        attention_mask = attention_mask.unsqueeze(0).to(self.device)
+
+        text_output = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            image_grid_thw =image_grid_thw,
+            output_hidden_states=True
+        )
+        # [max_sequence_length, batch, hidden_size] -> [batch, max_sequence_length, hidden_size]
+        # clone to have a contiguous tensor
+        prompt_embeds = text_output.hidden_states[-1].detach()
+        prompt_embeds = prompt_embeds[:,self.prompt_template_encode_start_idx: -self.prompt_template_encode_end_idx ,:]
+
+        text_ids = prepare_pos_ids(modality_id=0,
+                                   type='text',
+                                   start=(0, 0),
+                                   num_token=prompt_embeds.shape[1]).to(device, dtype=dtype)
+
+        return prompt_embeds, text_ids
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+        return latents
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1
+
+    def prepare_latents(
+        self,
+        image,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        prompt_embeds_length,
+        device,
+        generator,
+        latents=None,
+    ): 
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        image_latents, image_latents_ids = None, None
+
+        if image is not None:
+            image = image.to(device=self.device, dtype=dtype)
+            image_latents = self.vae.encode(image).latent_dist
+            image_latents = image_latents.mode()
+            image_latents = (image_latents - self.vae.config.shift_factor) *  self.vae.config.scaling_factor
+            image_latents = image_latents.to(device=self.device, dtype=dtype)
+            image_latents = self._pack_latents(
+                image_latents, batch_size, num_channels_latents, height, width
+            )
+
+            image_latents_ids = prepare_pos_ids(modality_id=2,
+                                           type='image',
+                                           start=(prompt_embeds_length,
+                                                  prompt_embeds_length),
+                                           height=height//2,
+                                           width=width//2).to(device, dtype=torch.float64)
+        
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+        image_latents = image_latents.to(device=device, dtype=dtype)
+        latents_ids = prepare_pos_ids(modality_id=1,
+                                        type='image',
+                                        start=(prompt_embeds_length,
+                                               prompt_embeds_length),
+                                        height=height//2,
+                                        width=width//2).to(device, dtype=torch.float64)
+
+        return latents, image_latents, latents_ids, image_latents_ids
+
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Optional[PipelineImageInput] = None,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        num_inference_steps: int = 50,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator,List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+
+        image_size = image[0].size if isinstance(image, list) else image.size
+        calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0]*1.0/image_size[1])
+        
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        
+        device = self._execution_device
+
+        image = self.image_processor.resize(image, calculated_height, calculated_width)
+        prompt_image = self.image_processor.resize(image, calculated_height//2, calculated_width//2)
+        image = self.image_processor.preprocess(image, calculated_height, calculated_width)
+        
+        negative_prompt = '' if negative_prompt is None else negative_prompt
+        negative_prompt = [negative_prompt]*num_images_per_prompt
+        prompt = [prompt]*num_images_per_prompt
+
+        prompt_embeds, text_ids = self.encode_prompt(
+            image=prompt_image,
+            prompts=prompt,
+            device=device,
+            dtype=torch.float64
+        )
+        negative_prompt_embeds, negative_text_ids = self.encode_prompt(
+            image=prompt_image,
+            prompts=negative_prompt, 
+            device = device, 
+            dtype=torch.float64
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = 16
+        latents, image_latents, latents_ids, image_latents_ids = self.prepare_latents(
+            image,
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            calculated_height,
+            calculated_width,
+            prompt_embeds.dtype,
+            prompt_embeds.shape[1],
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        guidance = None
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0).to(device)
+        else:
+            prompt_embeds = prompt_embeds.to(device)
+
+        if image is not None:
+            latent_image_ids = torch.cat([latents_ids, image_latents_ids], dim=0)
+        else:
+            latent_image_ids = latents_ids
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+                
+                latent_model_input = latents
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1)
+
+                latent_model_input = torch.cat([latent_model_input] * 2) if self.do_classifier_free_guidance else latent_model_input
+
+                timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
+                
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    return_dict=False,
+                )[0]
+                noise_pred = noise_pred[:, :image_seq_len]
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, calculated_height, calculated_width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            if latents.dtype != self.vae.dtype:
+                latents = latents.to(dtype=self.vae.dtype)
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        # self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return LongCatImagePipelineOutput(images=image)

longcat_image/pipelines/pipeline_output.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import List, Union
+
+import numpy as np
+import PIL.Image
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class LongCatImagePipelineOutput(BaseOutput):
+    """
+    Output class for Stable Diffusion 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)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]

longcat_image/utils/__init__.py

@@ -0,0 +1,5 @@
+from .dist_utils import get_world_size, get_rank, get_local_rank
+from .log_buffer import LogBuffer
+from .model_utils import split_quotation, prepare_pos_ids, calculate_shift, \
+                        retrieve_timesteps, optimized_scale,pack_latents, unpack_latents, \
+                        calculate_shift, prepare_pos_ids, encode_prompt, encode_prompt_edit

longcat_image/utils/dist_utils.py

@@ -0,0 +1,33 @@
+
+import os
+
+import torch.distributed as dist
+
+def is_distributed():
+    return get_world_size() > 1
+
+
+def get_world_size():
+    if not dist.is_available():
+        return 1
+    return dist.get_world_size() if dist.is_initialized() else 1
+
+
+def get_rank():
+    if not dist.is_available():
+        return 0
+    return dist.get_rank() if dist.is_initialized() else 0
+
+
+def get_local_rank():
+    if not dist.is_available():
+        return 0
+    return int(os.getenv('LOCAL_RANK', 0)) if dist.is_initialized() else 0
+
+
+def is_master():
+    return get_rank() == 0
+
+
+def is_local_master():
+    return get_local_rank() == 0

longcat_image/utils/log_buffer.py

@@ -0,0 +1,41 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from collections import OrderedDict
+
+import numpy as np
+
+
+class LogBuffer:
+
+    def __init__(self):
+        self.val_history = OrderedDict()
+        self.n_history = OrderedDict()
+        self.output = OrderedDict()
+        self.ready = False
+
+    def clear(self) -> None:
+        self.val_history.clear()
+        self.n_history.clear()
+        self.clear_output()
+
+    def clear_output(self) -> None:
+        self.output.clear()
+        self.ready = False
+
+    def update(self, vars: dict, count: int = 1) -> None:
+        assert isinstance(vars, dict)
+        for key, var in vars.items():
+            if key not in self.val_history:
+                self.val_history[key] = []
+                self.n_history[key] = []
+            self.val_history[key].append(var)
+            self.n_history[key].append(count)
+
+    def average(self, n: int = 0) -> None:
+        """Average latest n values or all values."""
+        assert n >= 0
+        for key in self.val_history:
+            values = np.array(self.val_history[key][-n:])
+            nums = np.array(self.n_history[key][-n:])
+            avg = np.sum(values * nums) / np.sum(nums)
+            self.output[key] = avg
+        self.ready = True

longcat_image/utils/model_utils.py

@@ -0,0 +1,293 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import re
+import math
+from PIL import Image
+
+import torch
+import torch.nn as nn
+import inspect
+from transformers import AutoTokenizer, AutoProcessor
+
+
+def pack_latents(latents, batch_size, num_channels_latents, height, width):
+    latents = latents.view(
+        batch_size, num_channels_latents, height // 2, 2, width // 2, 2
+    )
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+    latents = latents.reshape(
+        batch_size, (height // 2) * (width // 2), num_channels_latents * 4
+    )
+
+    return latents
+
+
+def unpack_latents(latents, height, width, vae_scale_factor):
+    batch_size, num_patches, channels = latents.shape
+
+    height = height // vae_scale_factor
+    width = width // vae_scale_factor
+
+    latents = latents.view(batch_size, height, width, channels // 4, 2, 2)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+    latents = latents.reshape(batch_size, channels //
+                              (2 * 2), height * 2, width * 2)
+
+    return latents
+
+
+def split_quotation(prompt, quote_pairs=None):
+    """
+    Implement a regex-based string splitting algorithm that identifies delimiters defined by single or double quote pairs. 
+
+    Examples::
+        >>> prompt_en = "Please write 'Hello' on the blackboard for me."
+        >>> print(split_quotation(prompt_en))
+        >>> # output: [('Please write ', False), ("'Hello'", True), (' on the blackboard for me.', False)]
+    """
+    word_internal_quote_pattern = re.compile(r"[a-zA-Z]+'[a-zA-Z]+")
+    matches_word_internal_quote_pattern = word_internal_quote_pattern.findall(prompt)
+    mapping_word_internal_quote = []
+
+    for i, word_src in enumerate(set(matches_word_internal_quote_pattern)):
+        word_tgt = 'longcat_$##$_longcat' * (i+1)
+        prompt = prompt.replace(word_src, word_tgt)
+        mapping_word_internal_quote.append([word_src, word_tgt])
+
+    if quote_pairs is None:
+        quote_pairs = [("'", "'"), ('"', '"'), ('‘', '’'), ('“', '”')]
+        quotes = ["'", '"', '‘', '’', '“', '”']
+        for q1 in quotes:
+            for q2 in quotes:
+                if (q1, q2) not in quote_pairs:
+                    quote_pairs.append((q1, q2))
+
+    pattern = '|'.join([re.escape(q1) + r'[^' + re.escape(q1+q2) +
+                       r']*?' + re.escape(q2) for q1, q2 in quote_pairs])
+
+    parts = re.split(f'({pattern})', prompt)
+
+    result = []
+    for part in parts:
+        for word_src, word_tgt in mapping_word_internal_quote:
+            part = part.replace(word_tgt, word_src)
+        if re.match(pattern, part):
+            if len(part):
+                result.append((part, True))
+        else:
+            if len(part):
+                result.append((part, False))
+    return result
+
+
+def encode_prompt(prompt: str, tokenizer: AutoTokenizer, text_tokenizer_max_length: int, prompt_template_encode_prefix: str, prompt_template_encode_suffix: str):
+
+    all_tokens = []
+    for clean_prompt_sub, matched in split_quotation(prompt):
+        if matched:
+            for sub_word in clean_prompt_sub:
+                tokens = tokenizer(sub_word, add_special_tokens=False)['input_ids']
+                all_tokens.extend(tokens)
+        else:
+            tokens = tokenizer(clean_prompt_sub, add_special_tokens=False)['input_ids']
+            all_tokens.extend(tokens)
+
+    all_tokens = all_tokens[:text_tokenizer_max_length]
+    text_tokens_and_mask = tokenizer.pad(
+        {'input_ids': [all_tokens]},
+        max_length=text_tokenizer_max_length,
+        padding='max_length',
+        return_attention_mask=True,
+        return_tensors='pt')
+    
+    prefix_tokens = tokenizer(prompt_template_encode_prefix, add_special_tokens=False)['input_ids']
+    suffix_tokens = tokenizer(prompt_template_encode_suffix, add_special_tokens=False)['input_ids']
+    prefix_tokens_mask = torch.tensor( [1]*len(prefix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype )
+    suffix_tokens_mask = torch.tensor( [1]*len(suffix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype )
+
+    prefix_tokens = torch.tensor(prefix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+    suffix_tokens = torch.tensor(suffix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+    
+    input_ids = torch.cat( (prefix_tokens, text_tokens_and_mask.input_ids[0], suffix_tokens), dim=-1 )
+    attention_mask = torch.cat( (prefix_tokens_mask, text_tokens_and_mask.attention_mask[0], suffix_tokens_mask), dim=-1 )
+
+
+    input_ids = text_tokens_and_mask['input_ids'].squeeze(0)
+    attention_mask = text_tokens_and_mask['attention_mask'].squeeze(0)
+
+    return input_ids, attention_mask
+
+def encode_prompt_edit(prompt: str, img: Image.Image,tokenizer: AutoTokenizer, image_processor_vl: AutoProcessor, text_tokenizer_max_length: int, prompt_template_encode_prefix: str, prompt_template_encode_suffix: str):
+    raw_vl_input = image_processor_vl(images=img,return_tensors="pt")
+    pixel_values = raw_vl_input['pixel_values'].squeeze(0)
+    image_grid_thw = raw_vl_input['image_grid_thw'].squeeze(0)
+
+    all_tokens = []
+    for clean_prompt_sub, matched in split_quotation(prompt):
+        if matched:
+            for sub_word in clean_prompt_sub:
+                tokens = tokenizer(sub_word, add_special_tokens=False)['input_ids']
+                all_tokens.extend(tokens)
+        else:
+            tokens = tokenizer(clean_prompt_sub, add_special_tokens=False)['input_ids']
+            all_tokens.extend(tokens)
+
+    all_tokens = all_tokens[:text_tokenizer_max_length]
+    text_tokens_and_mask = tokenizer.pad(
+        {'input_ids': [all_tokens]},
+        max_length=text_tokenizer_max_length,
+        padding='max_length',
+        return_attention_mask=True,
+        return_tensors='pt')
+    
+    text = prompt_template_encode_prefix
+    merge_length = image_processor_vl.merge_size**2
+    while  "<|image_pad|>" in text:
+        num_image_tokens = image_grid_thw.prod() // merge_length
+        text = text.replace( "<|image_pad|>", "<|placeholder|>" * num_image_tokens, 1)
+    text = text.replace("<|placeholder|>",  "<|image_pad|>")
+    
+    prefix_tokens = tokenizer(text, add_special_tokens=False)['input_ids']
+    suffix_tokens = tokenizer(prompt_template_encode_suffix, add_special_tokens=False)['input_ids']
+    prefix_tokens_mask = torch.tensor( [1]*len(prefix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype )
+    suffix_tokens_mask = torch.tensor( [1]*len(suffix_tokens), dtype = text_tokens_and_mask.attention_mask[0].dtype )
+
+    prefix_tokens = torch.tensor(prefix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+    suffix_tokens = torch.tensor(suffix_tokens,dtype=text_tokens_and_mask.input_ids.dtype)
+    
+    input_ids = torch.cat( (prefix_tokens, text_tokens_and_mask.input_ids[0], suffix_tokens), dim=-1 )
+    attention_mask = torch.cat( (prefix_tokens_mask, text_tokens_and_mask.attention_mask[0], suffix_tokens_mask), dim=-1 )
+
+
+    # input_ids = text_tokens_and_mask['input_ids'].squeeze(0)
+    # attention_mask = text_tokens_and_mask['attention_mask'].squeeze(0)
+
+    return input_ids, attention_mask, pixel_values, image_grid_thw
+
+
+def prepare_pos_ids(
+        modality_id=0,
+        type='text',
+        start=(0, 0),
+        num_token=None,
+        height=None,
+        width=None):
+    if type == 'text':
+        assert num_token
+        if height or width:
+            print(
+                'Warning: The parameters of height and width will be ignored in "text" type.')
+        pos_ids = torch.zeros(num_token, 3)
+        pos_ids[..., 0] = modality_id
+        pos_ids[..., 1] = torch.arange(num_token) + start[0]
+        pos_ids[..., 2] = torch.arange(num_token) + start[1]
+    elif type == 'image':
+        assert height and width
+        if num_token:
+            print('Warning: The parameter of num_token will be ignored in "image" type.')
+        pos_ids = torch.zeros(height, width, 3)
+        pos_ids[..., 0] = modality_id
+        pos_ids[..., 1] = (
+            pos_ids[..., 1] + torch.arange(height)[:, None] + start[0]
+        )
+        pos_ids[..., 2] = (
+            pos_ids[..., 2] + torch.arange(width)[None, :] + start[1]
+        )
+        pos_ids = pos_ids.reshape(height*width, 3)
+    else:
+        raise KeyError(f'Unknow type {type}, only support "text" or "image".')
+    # pos_ids = pos_ids[None, :].repeat(batch_size, 1, 1)
+    return pos_ids
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(
+            scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+# @torch.cuda.amp.autocast(dtype=torch.float32)
+@torch.amp.autocast('cuda', dtype=torch.float32)
+def optimized_scale(positive_flat, negative_flat):
+
+    # Calculate dot production
+    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
+
+    # Squared norm of uncondition
+    squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
+
+    # st_star = v_cond^T * v_uncond / ||v_uncond||^2
+    st_star = dot_product / squared_norm
+
+    return st_star

misc/accelerate_config.yaml

@@ -0,0 +1,21 @@
+compute_environment: LOCAL_MACHINE
+deepspeed_config:
+  deepspeed_multinode_launcher: standard
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  zero3_init_flag: true
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+use_cpu: false

misc/prompt_rewrite_api.py

@@ -0,0 +1,113 @@
+import os
+from openai import OpenAI
+import re
+
+SYSTEM_PROMPT_EN = """
+You are a prompt engineering expert for text-to-image models. Since text-to-image models have limited capabilities in understanding user prompts, you need to identify the core theme and intent of the user's input and improve the model's understanding accuracy and generation quality through optimization and rewriting. The rewrite must strictly retain all information from the user's original prompt without deleting or distorting any details.
+Specific requirements are as follows:
+1. The rewrite must not affect any information expressed in the user's original prompt; the rewritten prompt should use coherent natural language, avoid low-information redundant descriptions, and keep the rewritten prompt length as concise as possible.
+2. Ensure consistency between input and output languages: Chinese input yields Chinese output, and English input yields English output. The rewritten token count should not exceed 512.
+3. The rewritten description should further refine subject characteristics and aesthetic techniques appearing in the original prompt, such as lighting and textures.
+4. If the original prompt does not specify an image style, ensure the rewritten prompt uses a **realistic photography style**. If the user specifies a style, retain the user's style.
+5. When the original prompt requires reasoning to clarify user intent, use logical reasoning based on world knowledge to convert vague abstract descriptions into specific tangible objects (e.g., convert "the tallest animal" to "a giraffe").
+6. When the original prompt requires text generation, please use double quotes to enclose the text part (e.g., `"50% OFF"`).
+7. When the original prompt requires generating text-heavy scenes like webpages, logos, UIs, or posters, and no specific text content is specified, you need to infer appropriate text content and enclose it in double quotes. For example, if the user inputs: "A tourism flyer with a grassland theme," it should be rewritten as: "A tourism flyer with the image title 'Grassland'."
+8. When negative words exist in the original prompt, ensure the rewritten prompt does not contain negative words. For example, "a lakeside without boats" should be rewritten such that the word "boat" does not appear at all.
+9. Except for text content explicitly requested by the user, **adding any extra text content is prohibited**.
+Here are examples of rewrites for different types of prompts:
+# Examples (Few-Shot Learning)
+  1. User Input: An animal with nine lives.
+    Rewrite Output: A cat bathed in soft sunlight, its fur soft and glossy. The background is a comfortable home environment with light from the window filtering through curtains, creating a warm light and shadow effect. The shot uses a medium distance perspective to highlight the cat's leisurely and stretched posture. Light cleverly hits the cat's face, emphasizing its spirited eyes and delicate whiskers, adding depth and affinity to the image.
+  2. User Input: Create an anime-style tourism flyer with a grassland theme.
+    Rewrite Output: In the lower right of the center, a short-haired girl sits sideways on a gray, irregularly shaped rock. She wears a white short-sleeved dress and brown flat shoes, holding a bunch of small white flowers in her left hand, smiling with her legs hanging naturally. The girl has dark brown shoulder-length hair with bangs covering her forehead, brown eyes, and a slightly open mouth. The rock surface has textures of varying depths. To the girl's left and front is lush grass, with long, yellow-green blades, some glowing golden in the sunlight. The grass extends into the distance, forming rolling green hills that fade in color as they recede. The sky occupies the upper half of the picture, pale blue dotted with a few fluffy white clouds. In the upper left corner, there is a line of text in italic, dark green font reading "Explore Nature's Peace". Colors are dominated by green, blue, and yellow, fluid lines, and distinct light and shadow contrast, creating a quiet and comfortable atmosphere.
+  3. User Input: A Christmas sale poster with a red background, promoting a Buy 1 Get 1 Free milk tea offer.
+    Rewrite Output: The poster features an overall red tone, embellished with white snowflake patterns on the top and left side. The upper right features a bunch of holly leaves with red berries and a pine cone. In the upper center, golden 3D text reads "Christmas Heartwarming Feedback" centered, along with red bold text "Buy 1 Get 1". Below, two transparent cups filled with bubble tea are placed side by side; the tea is light brown with dark brown pearls scattered at the bottom and middle. Below the cups, white snow piles up, decorated with pine branches, red berries, and pine cones. A blurry Christmas tree is faintly visible in the lower right corner. The image has high clarity, accurate text content, a unified design style, a prominent Christmas theme, and a reasonable layout, providing strong visual appeal.
+  4. User Input: A woman indoors shot in natural light, smiling with arms crossed, showing a relaxed and confident posture.
+    Rewrite Output: The image features a young Asian woman with long dark brown hair naturally falling over her shoulders, with some strands illuminated by light, showing a soft sheen. Her features are delicate, with long eyebrows, bright and spirited dark brown eyes looking directly at the camera, revealing peace and confidence. She has a high nose bridge, full lips with nude lipstick, and corners of the mouth slightly raised in a faint smile. Her skin is fair, with cheeks and collarbones illuminated by warm light, showing a healthy ruddiness. She wears a black spaghetti strap tank top revealing graceful collarbone lines, and a thin gold necklace with small beads and metal bars glinting in the light. Her outer layer is a beige knitted cardigan, soft in texture with visible knitting patterns on the sleeves. Her arms are crossed over her chest, hands covered by the cardigan sleeves, in a relaxed posture. The background is a pure dark brown without extra decoration, making the figure the absolute focus. The figure is located in the center of the frame. Light enters from the upper right, creating bright spots on her left cheek, neck, and collarbone, while the right side is slightly shadowed, creating a three-dimensional and soft tone. Image details are clear, showcasing skin texture, hair, and clothing materials well. Colors are dominated by warm tones, with the combination of beige and dark brown creating a warm and comfortable atmosphere. The overall style is natural, elegant, and artistic.
+  5. User Input: Create a series of images showing the growth process of an apple from seed to fruit. The series should include four stages: 1. Sowing, 2. Seedling growth, 3. Plant maturity, 4. Fruit harvesting.
+    Rewrite Output: A 4-panel exquisite illustration depicting the growth process of an apple, capturing each stage precisely and clearly. 1. "Sowing": A close-up shot of a hand gently placing a small apple seed into fertile dark soil, with visible soil texture and the seed's smooth surface. The background is a soft-focus garden dotted with green leaves and sunlight filtering through. 2. "Seedling Growth": A young apple sapling breaks through the soil, stretching tender green leaves toward the sky. The scene is set in a vibrant garden illuminated by warm golden light, highlighting the seedling's delicate structure. 3. "Plant Maturity": A mature apple tree, lush with branches and leaves, covered in tender green foliage and developing small apples. The background is a vibrant orchard under a clear blue sky, with dappled sunlight creating a peaceful atmosphere. 4. "Fruit Harvesting": A hand reaches into the tree to pick a ripe red apple, its smooth skin glistening in the sun. The scene shows the abundance of the orchard, with baskets of apples in the background, giving a sense of fulfillment. Each illustration uses a realistic style, focusing on details and harmonious colors to showcase the natural beauty and development of the apple's life cycle.
+  6. User Input: If 1 represents red, 2 represents green, 3 represents purple, and 4 represents yellow, please generate a four-color rainbow based on this rule. The color order from top to bottom is 3142.
+    Rewrite Output: The image consists of four horizontally arranged colored stripes, ordered from top to bottom as purple, red, yellow, and green. A white number is centered on each stripe. The top purple stripe features the number "3", the red stripe below it has the number "1", the yellow stripe further down has the number "4", and the bottom green stripe has the number "2". All numbers use a sans-serif font in pure white, forming a sharp contrast with the background colors to ensure good readability. The stripes have high color saturation and a slight texture. The overall layout is simple and clear, with distinct visual effects and no extra decorative elements, emphasizing the numerical information. The image is high definition, with accurate colors and a consistent style, offering strong visual appeal.
+  7. User Input: A stone tablet carved with "Guan Guan Ju Jiu, On the River Isle", natural light, background is a Chinese garden.
+    Rewrite Output: An ancient stone tablet carved with "Guan Guan Ju Jiu, On the River Isle", the surface covered with traces of time, the writing clear and deep. Natural light falls from above, softly illuminating every detail of the stone tablet and enhancing its sense of history. The background is an elegant Chinese garden featuring lush bamboo forests, winding paths, and quiet pools, creating a serene and distant atmosphere. The overall picture uses a realistic style with rich details and natural light and shadow effects, highlighting the cultural heritage of the stone tablet and the classical beauty of the garden.
+# Output Format
+Please directly output the rewritten and optimized Prompt content. Do not include any explanatory language or JSON formatting, and do not add opening or closing quotes yourself."""
+
+
+SYSTEM_PROMPT_ZH = """
+你是一名文生图模型的prompt engineering专家。由于文生图模型对用户prompt的理解能力有限，你需要识别用户输入的核心主题和意图，并通过优化改写提升模型的理解准确性和生成质量。改写必须严格保留用户原始prompt的所有信息，不得删减或曲解任何细节。
+具体要求如下：
+1. 改写不能影响用户原始prompt里表达的任何信息，改写后的prompt应该使用连贯的自然语言表达,不要出现低信息量的冗余描述，尽可能保持改写后prompt长度精简。
+2. 请确保输入和输出的语言类型一致，中文输入中文输出，英文输入英文输出，改写后的token数量不要超过512个;
+3. 改写后的描述应当进一步完善原始prompt中出现的主体特征、美学技巧，如打光、纹理等；
+4. 如果原始prompt没有指定图片风格时，确保改写后的prompt使用真实摄影风格，如果用户指定了图片风格，则保留用户风格；
+5. 当原始prompt需要推理才能明确用户意图时，根据世界知识进行适当逻辑推理，将模糊抽象描述转化为具体指向事物（例：将"最高的动物"转化为"一头长颈鹿"）。
+6. 当原始prompt需要生成文字时，请使用双引号圈定文字部分，例：`"限时5折"`）。
+7. 当原始prompt需要生成网页、logo、ui、海报等文字场景时，且没有指定具体的文字内容时，需要推断出合适的文字内容，并使用双引号圈定，如用户输入：一个旅游宣传单，以草原为主题。应该改写成：一个旅游宣传单，图片标题为“草原”。
+8. 当原始prompt中存在否定词时，需要确保改写后的prompt不存在否定词，如没有船的湖边，改写后的prompt不能出现船这个词汇。
+9. 除非用户指定生成品牌logo，否则不要增加额外的品牌logo.
+10. 除了用户明确要求书写的文字内容外，**禁止增加任何额外的文字内容**。
+以下是针对不同类型prompt改写的示例：
+
+# Examples (Few-Shot Learning)
+  1. 用户输入: 九条命的动物。
+    改写输出: 一只猫，被柔和的阳光笼罩着，毛发柔软而富有光泽。背景是一个舒适的家居环境，窗外的光线透过窗帘，形成温馨的光影效果。镜头采用中距离视角，突出猫悠闲舒展的姿态。光线巧妙地打在猫的脸部，强调它灵动的眼睛和精致的胡须，增加画面的层次感与亲和力。
+  2. 用户输入: 制作一个动画风格的旅游宣传单，以草原为主题。
+    改写输出: 画面中央偏右下角，一个短发女孩侧身坐在灰色的不规则形状岩石上，她穿着白色短袖连衣裙和棕色平底鞋，左手拿着一束白色小花，面带微笑，双腿自然垂下。女孩的头发为深棕色，齐肩短发，刘海覆盖额头，眼睛呈棕色，嘴巴微张。岩石表面有深浅不一的纹理。女孩的左侧和前方是茂盛的草地，草叶细长，呈黄绿色，部分草叶在阳光下泛着金色的光芒，仿佛被阳光照亮。草地向远处延伸，形成连绵起伏的绿色山丘，山丘的颜色由近及远逐渐变浅。天空占据了画面的上半部分，呈淡蓝色，点缀着几朵白色蓬松的云彩。画面的左上角有一行文字，文字内容是斜体、深绿色的“Explore Nature's Peace”。色彩以绿色、蓝色和黄色为主，线条流畅，光影明暗对比明显，营造出一种宁静、舒适的氛围。
+  3. 用户输入: 一张以红色为背景的圣诞节促销海报，主要宣传奶茶买一送一的优惠活动。
+    改写输出: 海报整体呈现红色调，上方和左侧点缀着白色雪花图案，右上方有一束冬青叶和红色浆果，以及一个松果。海报中央偏上位置，金色立体字样“圣诞节 暖心回馈”居中排列，和红色粗体字“买1送1”。海报下方，两个装满珍珠奶茶的透明杯子并排摆放，杯中奶茶呈浅棕色，底部和中间散布着深棕色珍珠。杯子下方，堆积着白色雪花，雪花上装饰着松枝、红色浆果和松果。右下角隐约可见一棵模糊的圣诞树。图片清晰度高，文字内容准确，整体设计风格统一，圣诞主题突出，排版布局合理，具有较强的视觉吸引力。
+  4. 用户输入: 一位女性在室内以自然光线拍摄，她面带微笑，双臂交叉，展现出轻松自信的姿态。
+    改写输出: 画面中是一位年轻的亚洲女性，她拥有深棕色的长发，发丝自然地垂落在双肩，部分发丝被光线照亮，呈现出柔和的光泽。她的五官精致，眉毛修长，眼睛明亮有神，瞳孔呈深棕色，眼神直视镜头，流露出平和与自信。鼻梁挺拔，嘴唇丰满，涂有裸色系唇膏，嘴角微微上扬，展现出浅浅的微笑。她的肤色白皙，脸颊和锁骨处被暖色调的光线照亮，呈现出健康的红润感。她穿着一件黑色的细吊带背心，肩带纤细，露出优美的锁骨线条。脖颈上佩戴着一条金色的细项链，项链由小珠子和几个细长的金属条组成，在光线下闪烁着光泽。她的外搭是一件米黄色的针织开衫，材质柔软，袖子部分有明显的针织纹理。她双臂交叉在胸前，双手被开衫的袖子覆盖，姿态放松。背景是纯粹的深棕色，没有多余的装饰，使得人物成为画面的绝对焦点。人物位于画面中央。光线从画面的右上方射入，在人物的左侧脸颊、脖颈和锁骨处形成明亮的光斑，右侧则略显阴影，营造出立体感和柔和的影调。图像细节清晰，人物的皮肤纹理、发丝以及衣物材质都得到了很好的展现。色彩以暖色调为主，米黄色和深棕色的搭配营造出温馨舒适的氛围。整体呈现出一种自然、优雅且富有亲和力的艺术风格。
+  5. 用户输入：创作一系列图片，展现苹果从种子到结果的生长过程。该系列图片应包含以下四个阶段：1. 播种，2. 幼苗生长，3. 植物成熟，4. 果实采摘。
+    改写输出：一个4宫格的精美插图，描绘苹果的生长过程，精确清晰地捕捉每个阶段。1.“播种”：特写镜头，一只手轻轻地将一颗小小的苹果种子放入肥沃的深色土壤中，土壤的纹理和种子光滑的表面清晰可见。背景是花园的柔焦画面，点缀着绿色的树叶和透过树叶洒下的阳光。2.“幼苗生长”：一棵幼小的苹果树苗破土而出，嫩绿的叶子向天空舒展。场景设定在一个生机勃勃的花园中，温暖的金光照亮了它。幼苗的纤细结构。3.“植物的成熟”：一棵成熟的苹果树，枝繁叶茂，挂满了嫩绿的叶子和正在萌发的小苹果。背景是一片生机勃勃的果园，湛蓝的天空下，斑驳的阳光营造出宁静祥和的氛围。4.“采摘果实”：一只手伸向树上，摘下一个成熟的红苹果，苹果光滑的果皮在阳光下闪闪发光。画面展现了果园的丰收景象，背景中摆放着一篮篮的苹果，给人一种圆满满足的感觉。每幅插图都采用写实风格，注重细节，色彩和谐，展现了苹果生命周期的自然之美和发展过程。
+  6. 用户输入： 如果1代表红色，2代表绿色，3代表紫色，4代表黄色，请按照此规则生成四色彩虹。它的颜色顺序从上到下是3142
+    改写输出：图片由四个水平排列的彩色条纹组成，从上到下依次为紫色、红色、黄色和绿色。每个条纹上都居中放置一个白色数字。最上方的紫色条纹上是数字“3”，其下方红色条纹上是数字“1”，再下方黄色条纹上是数字“4”，最下方的绿色条纹上是数字“2”。所有数字均采用无衬线字体，颜色为纯白色，与背景色形成鲜明对比，确保了良好的可读性。条纹的颜色饱和度高，且带有轻微的纹理感，整体排版简洁明了，视觉效果清晰，没有多余的装饰元素，强调了数字信息本身。图片整体清晰度高，色彩准确，风格一致，具有较强的视觉吸引力。
+  7. 用户输入：石碑上刻着“关关雎鸠，在河之洲”，自然光照，背景是中式园林
+    改写输出：一块古老的石碑上刻着“关关雎鸠，在河之洲”，石碑表面布满岁月的痕迹，字迹清晰而深刻。自然光线从上方洒下，柔和地照亮石碑的每一个细节，增强了其历史感。背景是一座典雅的中式园林，园林中有翠绿的竹林、蜿蜒的小径和静谧的水池，营造出一种宁静而悠远的氛围。整体画面采用写实风格，细节丰富，光影效果自然，突出了石碑的文化底蕴和园林的古典美。
+# 输出格式
+请直接输出改写优化后的 Prompt 内容，不要包含任何解释性语言或 JSON 格式，不要自行添加开头或结尾的引号。
+"""
+
+
+def contains_chinese(text):
+    pattern = re.compile(r'[\u4e00-\u9fff]')
+    if bool(pattern.search(text)):
+      return 'zh'
+    return 'en'
+
+def prompt_rewrite_deepseek( prompt ):
+    client = OpenAI(api_key=os.environ.get('DEEPSEEK_API_KEY'),
+        base_url="https://api.deepseek.com" )
+    
+    try:
+        language = contains_chinese(prompt)
+        if language == 'zh':
+            question = SYSTEM_PROMPT_ZH + f"\n用户输入为：{prompt}\n改写后的prompt为："
+        else:
+            question = SYSTEM_PROMPT_EN + f"\nUser Input: {prompt}\nRewritten prompt:"
+            
+        response = client.chat.completions.create(
+            model="deepseek-chat",
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant"},
+                {"role": "user", "content": question }
+            ],
+            temperature=0.7,
+            max_tokens=512,
+            stream=False
+        )
+        rewrite_prompt = response.choices[0].message.content
+    except Exception as e:
+        rewrite_prompt = prompt
+        print(f"发生错误: {e}")
+    return rewrite_prompt
+
+if __name__ == "__main__":
+    # 输入格式
+    import os
+    os.environ['DEEPSEEK_API_KEY'] = 'sk-xxxxxxxxxxxxxxxxxxxxxxxxx'
+    prompt = '一个年轻的亚洲美女'
+    rewrite_prompt = prompt_rewrite_deepseek(  prompt )
+
+    print(rewrite_prompt)
+

nodes.py

@@ -0,0 +1,291 @@
+import math
+import os
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+# Add the directory containing this file to sys.path to allow importing longcat_image
+# which is located in the same directory
+current_file_path = Path(__file__).resolve()
+node_dir = current_file_path.parent
+if str(node_dir) not in sys.path:
+    sys.path.insert(0, str(node_dir))
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import AutoProcessor
+
+# ComfyUI imports
+import folder_paths
+import comfy.sd
+import comfy.utils
+import comfy.model_management
+import comfy.samplers
+import comfy.model_patcher
+import comfy.clip_vision
+
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.pipelines import LongCatImageEditPipeline, LongCatImagePipeline
+from longcat_image.pipelines.pipeline_longcat_image_edit import calculate_dimensions
+from longcat_image.utils.model_utils import prepare_pos_ids, calculate_shift, retrieve_timesteps
+
+
+# Helpers for device/dtype
+def _select_device() -> torch.device:
+    return comfy.model_management.get_torch_device()
+
+def _select_dtype(device: torch.device) -> torch.dtype:
+    # ComfyUI usually handles this, but we can default to fp16 for inference
+    return torch.float16
+
+def _torch_image_to_pil_list(image: torch.Tensor) -> List[Image.Image]:
+    if image.dim() == 3:
+        image = image.unsqueeze(0)
+    image = image.clamp(0, 1).detach().cpu()
+    pil_images = []
+    for img in image:
+        np_img = (img.numpy() * 255.0).round().astype(np.uint8)
+        pil_images.append(Image.fromarray(np_img))
+    return pil_images
+
+def _pil_list_to_torch(images: List[Image.Image]) -> torch.Tensor:
+    tensors = []
+    for img in images:
+        np_img = np.array(img, copy=False)
+        if np_img.ndim == 2:
+            np_img = np.repeat(np_img[..., None], 3, axis=2)
+        tensor = torch.from_numpy(np_img.astype(np.float32) / 255.0)
+        tensors.append(tensor)
+    if not tensors:
+        return torch.zeros((1, 512, 512, 3), dtype=torch.float32)
+    return torch.stack(tensors, dim=0)
+
+def _round_to_multiple(value: int, multiple: int = 16) -> int:
+    return int(math.ceil(value / multiple) * multiple)
+
+
+# --- Loaders & Wrappers ---
+
+class LongCatCheckpointLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
+    FUNCTION = "load_checkpoint"
+    CATEGORY = "longcat"
+
+    def load_checkpoint(self, ckpt_name: str):
+        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+        
+        # Load state dict
+        sd = comfy.utils.load_torch_file(ckpt_path)
+        
+        # Load Model (Transformer)
+        # We assume the checkpoint contains the transformer weights or is a diffusers dump
+        # If it's a single file, we need to know the prefix or structure.
+        # For now, let's assume it's a standard Diffusers-style checkpoint loaded as a dict.
+        # If it's a folder, comfy.utils.load_torch_file might fail or return dict of files?
+        # folder_paths usually returns file paths.
+        
+        # Initialize Transformer
+        # We need config. If not in checkpoint, we might need default config.
+        # Assuming standard LongCat config for now.
+        transformer = LongCatImageTransformer2DModel(
+            patch_size=2,
+            in_channels=16,
+            num_layers=28,
+            num_single_layers=10, # Example default, adjust as needed
+            attention_head_dim=72,
+            num_attention_heads=16,
+            joint_attention_dim=768, # Example default
+            pooled_projection_dim=768, # Example default
+            axes_dims_rope=[16, 56, 56],
+        )
+        
+        # Load weights into transformer
+        # We might need to filter keys if the checkpoint contains CLIP/VAE too.
+        # This is a simplification. Real implementation needs robust key mapping.
+        # transformer.load_state_dict(sd, strict=False) 
+        
+        # Wrap in ModelPatcher
+        model = comfy.model_patcher.ModelPatcher(transformer, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
+        
+        # Load CLIP
+        # Using Comfy's CLIP loader logic if possible, or creating one.
+        # clip = comfy.sd.load_clip(ckpt_path) # This might work if it's a supported format
+        
+        # Load VAE
+        # vae = comfy.sd.load_vae(ckpt_path)
+        
+        # Placeholder for actual loading logic which depends on file format
+        # For this refactor, we assume we return standard objects.
+        
+        return (model, None, None) # Placeholder
+
+
+# --- Encoders ---
+
+class TextEncodeLongCatImage:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "clip": ("CLIP",),
+                "prompt": ("STRING", {"multiline": True}),
+                "enable_prompt_rewrite": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
+    CATEGORY = "longcat"
+
+    def encode(self, clip: Any, prompt: str, enable_prompt_rewrite: bool):
+        # clip is now a comfy.sd.CLIP object
+        # We can use standard encoding or custom logic
+        
+        tokens = clip.tokenize(prompt)
+        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
+        
+        # Pack into CONDITIONING
+        return ([[cond, {"pooled_output": pooled}]],)
+
+
+class TextEncodeLongCatImageEdit:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "clip": ("CLIP",),
+                "vae": ("VAE",),
+                "image1": ("IMAGE",),
+                "prompt": ("STRING", {"multiline": True}),
+            },
+            "optional": {
+                "image2": ("IMAGE",),
+                "image3": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
+    CATEGORY = "longcat"
+
+    def encode(self, clip: Any, vae: Any, image1: torch.Tensor, prompt: str, image2: Optional[torch.Tensor] = None, image3: Optional[torch.Tensor] = None):
+        # Custom logic for edit encoding
+        # We need to access the internal tokenizer/encoder of the CLIP object if possible
+        # or use standard encoding and add image tokens?
+        
+        # Placeholder for complex logic
+        return ([],)
+
+
+# --- VAE ---
+
+class VAEEncodeLongCat:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "vae": ("VAE",),
+                "image": ("IMAGE",),
+            }
+        }
+    
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "encode"
+    CATEGORY = "longcat"
+    
+    def encode(self, vae: Any, image: torch.Tensor):
+        return (vae.encode(image[:,:,:,:3]),)
+
+
+class VAEDecodeLongCat:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "vae": ("VAE",),
+                "samples": ("LATENT",),
+            }
+        }
+    
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "decode"
+    CATEGORY = "longcat"
+    
+    def decode(self, vae: Any, samples: Dict):
+        return (vae.decode(samples["samples"]),)
+
+
+# --- Sampler ---
+
+class LongCatSampler:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "latent_image": ("LATENT",),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 30, "min": 1, "max": 1000}),
+                "cfg": ("FLOAT", {"default": 4.5, "min": 0.0, "max": 100.0}),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "sample"
+    CATEGORY = "longcat"
+
+    def sample(self, model: Any, positive: List, negative: List, latent_image: Dict, seed: int, steps: int, cfg: float, sampler_name: str, scheduler: str, denoise: float):
+        # Use ComfyUI's common_ksampler logic
+        # This handles the sampling loop using the provided sampler/scheduler
+        
+        return comfy.samplers.common_ksampler(
+            model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise
+        )
+
+
+class LongCatImageSizeScale:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "target_area": ("INT", {"default": 1024 * 1024, "min": 256 * 256, "max": 4096 * 4096}),
+                "upscale_method": (["nearest-exact", "bilinear", "area", "bicubic", "lanczos"],),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "INT", "INT")
+    RETURN_NAMES = ("image", "width", "height")
+    FUNCTION = "scale"
+    CATEGORY = "longcat"
+
+    def scale(self, image: torch.Tensor, target_area: int, upscale_method: str):
+        if image.dim() == 3:
+            image = image.unsqueeze(0)
+        b, h, w, c = image.shape
+        ratio = w / float(h)
+        width = math.sqrt(target_area * ratio)
+        height = width / ratio
+        width = _round_to_multiple(int(round(width)), 16)
+        height = _round_to_multiple(int(round(height)), 16)
+        
+        # Use ComfyUI's common_upscale
+        scaled = comfy.utils.common_upscale(image.permute(0, 3, 1, 2), width, height, upscale_method, "center")
+        scaled = scaled.permute(0, 2, 3, 1)
+        
+        return (scaled, int(width), int(height))

requirements.txt

@@ -0,0 +1,8 @@
+torch>=2.1.0
+transformers>=4.41.0
+diffusers>=0.30.0
+accelerate>=0.27.0
+pillow>=10.0.0
+numpy>=1.24.0
+openai>=1.12.0
+pytest>=7.4.0

scripts/inference_edit.py

@@ -0,0 +1,37 @@
+import torch
+from PIL import Image
+from transformers import AutoProcessor
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.pipelines import LongCatImageEditPipeline
+
+if __name__ == '__main__':
+
+    device = torch.device('cuda')
+    checkpoint_dir = './weights/LongCat-Image-Edit'
+    text_processor = AutoProcessor.from_pretrained( checkpoint_dir, subfolder = 'tokenizer'  )
+    transformer = LongCatImageTransformer2DModel.from_pretrained( checkpoint_dir , subfolder = 'transformer', 
+        torch_dtype=torch.bfloat16, use_safetensors=True).to(device)
+
+    pipe = LongCatImageEditPipeline.from_pretrained(
+        checkpoint_dir,
+        transformer=transformer,
+        text_processor=text_processor,
+        torch_dtype=torch.bfloat16,
+    )
+    # pipe.to(device, torch.bfloat16)  # Uncomment for high VRAM devices (Faster inference)
+    pipe.enable_model_cpu_offload()  # Offload to CPU to save VRAM (Required ~19 GB); slower but prevents OOM
+
+    generator = torch.Generator("cpu").manual_seed(43)
+    img = Image.open('assets/test.png').convert('RGB')
+    prompt = '将猫变成狗'
+    image = pipe(
+        img,
+        prompt,
+        negative_prompt='',
+        guidance_scale=4.5,
+        num_inference_steps=50,
+        num_images_per_prompt=1,
+        generator=generator
+    ).images[0]
+
+    image.save('./edit_example.png')

scripts/inference_t2i.py

@@ -0,0 +1,44 @@
+import torch
+
+from transformers import AutoProcessor
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.pipelines import LongCatImagePipeline
+from misc.prompt_rewrite_api import prompt_rewrite_deepseek
+
+if __name__ == '__main__':
+
+    device = torch.device('cuda')
+    checkpoint_dir = './weights/LongCat-Image'
+
+    text_processor = AutoProcessor.from_pretrained( checkpoint_dir, subfolder = 'tokenizer'  )
+    transformer = LongCatImageTransformer2DModel.from_pretrained( checkpoint_dir , subfolder = 'transformer', 
+        torch_dtype=torch.bfloat16, use_safetensors=True).to(device)
+
+    pipe = LongCatImagePipeline.from_pretrained(
+        checkpoint_dir,
+        transformer=transformer,
+        text_processor=text_processor,
+        torch_dtype=torch.bfloat16
+    )
+    # pipe.to(device, torch.bfloat16)  # Uncomment for high VRAM devices (Faster inference)
+    pipe.enable_model_cpu_offload()  # Offload to CPU to save VRAM (Required ~17 GB); slower but prevents OOM
+
+    prompt = '一个年轻的亚裔女性，身穿黄色针织衫，搭配白色项链。她的双手放在膝盖上，表情恬静。背景是一堵粗糙的砖墙，午后的阳光温暖地洒在她身上，营造出一种宁静而温馨的氛围。镜头采用中距离视角，突出她的神态和服饰的细节。光线柔和地打在她的脸上，强调她的五官和饰品的质感，增加画面的层次感与亲和力。整个画面构图简洁，砖墙的纹理与阳光的光影效果相得益彰，突显出人物的优雅与从容。'
+
+    enable_prompt_rewrite_api = False
+    if enable_prompt_rewrite_api:
+        prompt = prompt_rewrite_deepseek(prompt)
+
+    image = pipe(
+        prompt,
+        negative_prompt='',
+        height=768,
+        width=1344,
+        guidance_scale=4.5,
+        num_inference_steps=50,
+        num_images_per_prompt=1,
+        generator= torch.Generator("cpu").manual_seed(43),
+        enable_cfg_renorm=True,
+        enable_prompt_rewrite=True if not enable_prompt_rewrite_api else False
+    ).images[0]
+    image.save('./t2i_example.png')

setup.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+from setuptools import find_packages, setup
+
+
+this_dir = Path(__file__).parent
+readme = (this_dir / "README.md")
+long_description = readme.read_text(encoding="utf-8") if readme.exists() else "LongCat ComfyUI nodes and pipelines."
+
+setup(
+    name="comfyui-longcat",
+    version="0.1.0",
+    description="ComfyUI nodes for LongCat image generation and editing pipelines",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="LongCat Team",
+    packages=find_packages(exclude=("tests", "train_examples")),
+    python_requires=">=3.9",
+    install_requires=[
+        "torch>=2.1.0",
+        "transformers>=4.41.0",
+        "diffusers>=0.30.0",
+        "accelerate>=0.27.0",
+        "pillow>=10.0.0",
+        "numpy>=1.24.0",
+        "openai>=1.12.0",
+    ],
+    extras_require={
+        "dev": ["pytest>=7.4.0"],
+    },
+    include_package_data=True,
+    package_data={
+        "": ["*.yaml"],
+    },
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+    ],
+)

tests/test_nodes.py

@@ -0,0 +1,66 @@
+import sys
+from unittest.mock import MagicMock
+
+# Mock ComfyUI modules BEFORE importing nodes
+sys.modules["folder_paths"] = MagicMock()
+sys.modules["comfy"] = MagicMock()
+sys.modules["comfy.sd"] = MagicMock()
+sys.modules["comfy.utils"] = MagicMock()
+sys.modules["comfy.model_management"] = MagicMock()
+sys.modules["comfy.samplers"] = MagicMock()
+sys.modules["comfy.model_patcher"] = MagicMock()
+sys.modules["comfy.clip_vision"] = MagicMock()
+
+# Setup mock values
+sys.modules["comfy.samplers"].KSampler.SAMPLERS = ["euler"]
+sys.modules["comfy.samplers"].KSampler.SCHEDULERS = ["normal"]
+
+import pytest
+import torch
+from pathlib import Path
+
+# Add repo root to sys.path
+current_file_path = Path(__file__).resolve()
+repo_root = current_file_path.parent.parent.parent
+if str(repo_root) not in sys.path:
+    sys.path.insert(0, str(repo_root))
+
+from comfyui_longcat.nodes import (
+    LongCatCheckpointLoader,
+    TextEncodeLongCatImage,
+    TextEncodeLongCatImageEdit,
+    VAEEncodeLongCat,
+    VAEDecodeLongCat,
+    LongCatSampler,
+    LongCatImageSizeScale,
+)
+
+def test_loader():
+    # Mock folder_paths
+    sys.modules["folder_paths"].get_full_path.return_value = "dummy_path.safetensors"
+    sys.modules["comfy.utils"].load_torch_file.return_value = {}
+    
+    loader = LongCatCheckpointLoader()
+    model, clip, vae = loader.load_checkpoint("dummy.safetensors")
+    
+    assert model is not None
+    # CLIP/VAE are None in placeholder implementation
+    
+def test_scale():
+    scaler = LongCatImageSizeScale()
+    image = torch.zeros((1, 100, 100, 3))
+    
+    # Mock common_upscale
+    sys.modules["comfy.utils"].common_upscale.return_value = torch.zeros((1, 3, 1024, 1024))
+    
+    scaled, w, h = scaler.scale(image, 1024*1024, "bicubic")
+    assert w == 1024
+    assert h == 1024
+
+def test_sampler():
+    sampler = LongCatSampler()
+    # Just check if it calls common_ksampler
+    sys.modules["comfy.samplers"].common_ksampler.return_value = ({"samples": torch.zeros(1)},)
+    
+    res = sampler.sample(MagicMock(), [], [], {}, 0, 1, 1.0, "euler", "normal", 1.0)
+    assert res is not None

train_examples/README.md

@@ -0,0 +1,101 @@
+
+## Training Example
+
+### 1. training example for sft or lora
+
+- Data format
+
+You need to create a jsonl file with key-values in the table below:
+
+|    key_word       | Required  |  Description     |   Example   |
+|:---------------:| :------:  |:----------------:|:-----------:|
+|   `img_path`  | Required  |  image path               |     `./data_example/images/0.png`        | 
+|   `prompt`     | Required  |    text               |  `A lovely little girl.` | 
+|   `width`       | Required  |    image width    |     ` 1024 `       | 
+|   `height`      | Required  |    image height   |     ` 1024 `       | 
+
+
+
+- Tainging Scripts
+
+```bash
+
+bash ./train_examples/sft/train.sh
+
+# All training setting in train_config.yaml
+
+# --data_csv_root:  data csv_filepath
+# --aspect_ratio_type: data bucketing strategy, mar_256、mar_512、mar_1024
+# --pretrained_model_name_or_path: root directory of the model
+# --diffusion_pretrain_weight: if a specified diffusion weight path is provided, load the model parameters from the current directory.
+# --work_dir: the save root directory for ckpt and logs
+# --resume_from_checkpoint: If 'resume_from_checkpoint' is set to 'latest', load the most recent step checkpoint. If a specific directory is provided, resume training from that directory.
+
+```
+
+
+### 2. training example for dpo
+
+- Data format
+
+You need to create a txt file with key-values in the table below:
+
+
+|    key_word       | Required  |  Description     |   Example   |
+|:---------------:| :------:  |:----------------:|:-----------:|
+|   `img_path_win`  | Required  |  win image path               |     `./data_example/images/0.png`        | 
+|   `img_path_lose`  | Required  |  lose image path               |     `./data_example/images/1.png`        | 
+|   `prompt`     | Required  |    text               |  `A lovely little girl.` | 
+|   `width`       | Required  |    image width    |     ` 1024 `       | 
+|   `height`      | Required  |    image height   |     ` 1024 `       | 
+
+
+- Tainging Scripts
+
+```bash
+
+bash ./train_examples/dpo/train.sh
+
+# All training setting in train_config.yaml
+
+# --data_txt_root:  data txt_filepath
+# --aspect_ratio_type: data bucketing strategy, mar_256、mar_512、mar_1024
+# --pretrained_model_name_or_path: root directory of the model
+# --diffusion_pretrain_weight: if a specified diffusion weight path is provided, load the model parameters from the current directory.
+# --work_dir: the save root directory for ckpt and logs
+# --resume_from_checkpoint: If 'resume_from_checkpoint' is set to 'latest', load the most recent step checkpoint. If a specific directory is provided, resume training from that directory.
+
+```
+
+### 3. training example for image-edit
+
+- Data format
+
+You need to create a txt file with key-values in the table below:
+
+
+|    key_word       | Required  |  Description     |   Example   |
+|:---------------:| :------:  |:----------------:|:-----------:|
+|   `img_path`  | Required  |  edited image path      |     `./data_example/images/0_edited.png`        | 
+|   `ref_img_path`  | Required  |  raw image path     |     `./data_example/images/0.png`        | 
+|   `prompt`     | Required  |    edit instruction     |  `change the dog to cat.` | 
+|   `width`       | Required  |    image width    |     ` 1024 `       | 
+|   `height`      | Required  |    image height   |     ` 1024 `       | 
+
+
+- Tainging Scripts
+
+```bash
+
+bash ./train_examples/edit/train.sh
+
+# All training setting in train_config.yaml
+
+# --data_txt_root:  data txt_filepath
+# --aspect_ratio_type: data bucketing strategy, mar_256、mar_512、mar_1024
+# --pretrained_model_name_or_path: root directory of the model
+# --diffusion_pretrain_weight: if a specified diffusion weight path is provided, load the model parameters from the current directory.
+# --work_dir: the save root directory for ckpt and logs
+# --resume_from_checkpoint: If 'resume_from_checkpoint' is set to 'latest', load the most recent step checkpoint. If a specific directory is provided, resume training from that directory.
+
+```

train_examples/dpo/train.sh

@@ -0,0 +1,13 @@
+export TOKENIZERS_PARALLELISM=False
+export NCCL_DEBUG=INFO
+export NCCL_TIMEOUT=12000
+
+script_dir=$(cd -- "$(dirname -- "$0")" &> /dev/null && pwd -P)
+project_root=$(dirname "$(dirname "$script_dir")")
+echo "script_dir" ${script_dir}
+
+deepspeed_config_file=${project_root}/misc/accelerate_config.yaml
+
+accelerate launch  --mixed_precision bf16 --num_processes 8 --config_file ${deepspeed_config_file} \
+${script_dir}/train_dpo.py \
+--config ${script_dir}/train_config.yaml

train_examples/dpo/train_config.yaml

@@ -0,0 +1,41 @@
+# 1.Data setting
+data_txt_root: '/dataset/example/train_data_info.txt'  # data csv_filepath
+resolution: 1024 
+aspect_ratio_type: 'mar_1024'  # data bucketing strategy, mar_256、mar_512、mar_1024
+null_text_ratio: 0.1
+dataloader_num_workers: 8
+train_batch_size: 4
+
+# 2. Model setting
+text_tokenizer_max_length: 512  # tokenizer max len 
+pretrained_model_name_or_path: "/xxx/weights/Longcat-Image-Dev" # root directory of the model，with vae、transformer、scheduler eta;
+diffusion_pretrain_weight:  null # if a specified diffusion weight path is provided, load the model parameters from the current directory.
+use_dynamic_shifting: true #  scheduler dynamic shifting
+resume_from_checkpoint: latest
+  # - "latest"  # Loads most recent step checkpoint
+  # - "/path/to/checkpoint"  # Resumes from specified directory
+
+# 3. Training setting
+use_ema: False
+ema_rate: 0.999
+mixed_precision: 'bf16'
+max_train_steps: 100000
+gradient_accumulation_steps: 1
+gradient_checkpointing: true
+gradient_clip: 1.0
+learning_rate: 5.0e-6
+adam_weight_decay: 1.0e-3
+adam_epsilon: 1.0e-8
+adam_beta1: 0.9
+adam_beta2: 0.999
+lr_num_cycles: 1
+lr_power: 1.0
+lr_scheduler: 'constant'
+lr_warmup_steps: 1000
+beta_dpo: 2000
+
+#4. Log setting
+log_interval: 20
+save_model_steps: 1000
+work_dir: 'output/sft_model'
+seed: 43

train_examples/dpo/train_dataset.py

@@ -0,0 +1,190 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+import os
+import random
+import traceback
+import math
+import json
+import numpy as np
+
+import torch
+import torch.distributed as dist
+import torchvision.transforms as T
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoTokenizer
+from PIL import Image
+from tqdm import tqdm
+
+from longcat_image.dataset import MULTI_RESOLUTION_MAP
+from longcat_image.utils import encode_prompt
+from longcat_image.dataset import MultiResolutionDistributedSampler
+
+Image.MAX_IMAGE_PIXELS = 2000000000
+
+MAX_RETRY_NUMS = 100
+
+class DpoPairDataSet(torch.utils.data.Dataset):
+    def __init__(self,
+                 cfg: dict,
+                 txt_root: str,
+                 tokenizer: AutoTokenizer,
+                 resolution: tuple = (1024, 1024)):
+        super(DpoPairDataSet, self).__init__()
+        self.resolution = resolution
+        self.text_tokenizer_max_length = cfg.text_tokenizer_max_length
+        self.null_text_ratio = cfg.null_text_ratio
+        self.aspect_ratio_type = cfg.aspect_ratio_type
+        self.aspect_ratio = MULTI_RESOLUTION_MAP[self.aspect_ratio_type]
+        self.tokenizer = tokenizer
+
+        self.total_datas = []
+        self.data_resolution_infos = []
+        with open(txt_root, 'r') as f:
+            lines = f.readlines()
+            for line in tqdm(lines):
+                data = json.loads(line.strip())
+                try:
+                    height, widht = int(data['height']), int(data['width'])
+                    self.data_resolution_infos.append((height, widht))
+                    self.total_datas.append(data)
+                except Exception as e:
+                    print(f'get error {e}, data {data}.')
+                    continue
+        self.data_nums = len(self.total_datas)
+
+    def transform_img(self, image, original_size, target_size):
+        img_h, img_w = original_size
+        target_height, target_width = target_size
+
+        original_aspect = img_h / img_w  # height/width
+        crop_aspect = target_height / target_width
+
+        if original_aspect >= crop_aspect:
+            resize_width = target_width
+            resize_height = math.ceil(img_h * (target_width/img_w))
+        else:
+            resize_width = math.ceil(img_w * (target_height/img_h))
+            resize_height = target_height
+
+        image = T.Compose([
+            T.Resize((resize_height, resize_width),interpolation=InterpolationMode.BICUBIC),  # Image.LANCZOS
+            T.CenterCrop((target_height, target_width)),
+            T.ToTensor(),
+            T.Normalize([.5], [.5]),
+        ])(image)
+
+        return image
+
+    def __getitem__(self, index_tuple):
+        index, target_size = index_tuple
+
+        for _ in range(MAX_RETRY_NUMS):
+            try:
+                item = self.total_datas[index]
+                img_path_win = item["img_path_win"]
+                img_path_lose = item["img_path_lose"]
+
+                prompt = item['prompt']
+
+                if random.random() < self.null_text_ratio:
+                    prompt = ''
+
+                raw_image_win = Image.open(img_path_win).convert('RGB')
+                raw_image_lose = Image.open(img_path_lose).convert('RGB')
+                assert raw_image_win is not None and raw_image_lose is not None
+                img_w, img_h = raw_image_win.size
+
+                raw_image_win = self.transform_img(raw_image_win, original_size=(
+                    img_h, img_w), target_size= target_size )
+                raw_image_lose = self.transform_img(raw_image_lose, original_size=(
+                    img_h, img_w), target_size= target_size )
+
+                input_ids,attention_mask = encode_prompt(prompt, self.tokenizer, self.text_tokenizer_max_length)
+                
+                return {"image_win": raw_image_win, "image_lose": raw_image_lose,  "prompt": prompt, 'input_ids': input_ids, 'attention_mask': attention_mask}
+
+            except Exception as e:
+                traceback.print_exc()
+                print(f"failed read data {e}!!!")
+                index = random.randint(0, self.data_nums-1)
+
+    def __len__(self):
+        return self.data_nums
+
+    def collate_fn(self, batchs):
+        images_win = torch.stack([example["image_win"] for example in batchs])
+        images_lose = torch.stack([example["image_lose"] for example in batchs])
+        input_ids = torch.stack([example["input_ids"] for example in batchs])
+        attention_mask = torch.stack([example["attention_mask"] for example in batchs])
+        prompts = [example['prompt'] for example in batchs]
+        batch_dict = {
+            "images_win": images_win,
+            "images_lose": images_lose,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "prompts": prompts,
+        }
+        return batch_dict
+
+
+def build_dataloader(cfg: dict,
+                     csv_root: str,
+                     tokenizer: AutoTokenizer,
+                     resolution: tuple = (1024, 1024)):
+    dataset = DpoPairDataSet(cfg, csv_root, tokenizer, resolution)
+
+    sampler = MultiResolutionDistributedSampler(batch_size=cfg.train_batch_size, dataset=dataset,
+                                                data_resolution_infos=dataset.data_resolution_infos,
+                                                bucket_info=dataset.aspect_ratio,
+                                                epoch=0,
+                                                num_replicas=None,
+                                                rank=None
+                                                )
+
+    train_loader = torch.utils.data.DataLoader(
+        dataset,
+        collate_fn=dataset.collate_fn,
+        batch_size=cfg.train_batch_size,
+        num_workers=cfg.dataloader_num_workers,
+        sampler=sampler,
+        shuffle=None,
+    )
+    return train_loader
+
+
+if __name__ == '__main__':
+    import sys
+    import argparse
+    from torchvision.transforms.functional import to_pil_image
+
+    txt_root = 'xxxx'
+
+    cfg = argparse.Namespace(
+        txt_root=txt_root,
+        text_tokenizer_max_length=256,
+        resolution=1024,
+        text_encoder_path="xxx",
+        center_crop=True,
+        dataloader_num_workers=0,
+        null_text_ratio=0.1,
+        train_batch_size=16,
+        seed=0,
+        aspect_ratio_type='mar_1024',
+        revision=None)
+
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(cfg.text_encoder_path, trust_remote_code=True)
+    data_loader = build_dataloader(cfg, cfg.csv_root, tokenizer, cfg.resolution)
+
+    _oroot = f'./debug_data_example_show'
+    os.makedirs(_oroot, exist_ok=True)
+
+    cnt = 0
+    for epoch in range(1):
+        print(f"Start, epoch {epoch}!!!")
+        for i_batch, batch in enumerate(data_loader):
+            print(batch['attention_mask'].shape)
+            print(batch['images_win'].shape,'-',batch['images_lose'].shape,)
+            
+            if cnt > 100:
+                break
+            cnt += 1

train_examples/dpo/train_dpo.py

@@ -0,0 +1,465 @@
+import os
+import sys
+import time
+import warnings
+
+import argparse
+import yaml
+import torch
+import math
+import logging
+import transformers
+import diffusers
+import torch.nn.functional as F
+
+from pathlib import Path
+from transformers import Qwen2Model, Qwen2TokenizerFast
+from accelerate import Accelerator, InitProcessGroupKwargs
+from accelerate.utils import ProjectConfiguration, set_seed
+from accelerate.logging import get_logger
+from diffusers.models import AutoencoderKL
+from diffusers.optimization import get_scheduler
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.training_utils import EMAModel
+from diffusers.utils.import_utils import is_xformers_available
+from transformers import AutoTokenizer, AutoModel
+from train_dataset import build_dataloader
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.utils import LogBuffer
+from longcat_image.utils import pack_latents, unpack_latents, calculate_shift, prepare_pos_ids
+
+warnings.filterwarnings("ignore")  # ignore warning
+
+current_file_path = Path(__file__).resolve()
+sys.path.insert(0, str(current_file_path.parent.parent))
+
+logger = get_logger(__name__)
+
+def train(global_step=0):
+
+    # Train!
+    total_batch_size = args.train_batch_size * \
+        accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataloader.dataset)}")
+    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    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}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+
+    last_tic = time.time()
+
+    # Now you train the model
+    for epoch in range(first_epoch, args.num_train_epochs + 1):
+        data_time_start = time.time()
+        data_time_all = 0
+
+        for step, batch in enumerate(train_dataloader):
+            images_win = batch['images_win']
+            images_lose = batch['images_lose']
+            half_batch_size = images_win.shape[0]
+            image = torch.concat([images_win, images_lose], dim=0)
+            
+            data_time_all += time.time() - data_time_start
+
+            with torch.no_grad():
+                latents = vae.encode(image.to(weight_dtype).to(accelerator.device)).latent_dist.sample()
+                latents = latents.to(dtype=(weight_dtype))
+                latents = (latents - vae.config.shift_factor) * vae.config.scaling_factor
+
+            text_input_ids = batch['input_ids'].to(accelerator.device)
+            text_attention_mask = batch['attention_mask'].to(accelerator.device)
+
+            with torch.no_grad():
+                text_output = text_encoder(
+                    input_ids=text_input_ids,
+                    attention_mask=text_attention_mask,
+                    output_hidden_states=True
+                )
+                prompt_embeds = text_output.hidden_states[-1].clone().detach()
+
+            prompt_embeds = prompt_embeds.to(weight_dtype)
+            prompt_embeds = torch.concat([prompt_embeds, prompt_embeds], dim=0) # [batch_size, 256, 4096]
+            prompt_embeds = prompt_embeds[:,args.prompt_template_encode_start_idx: -args.prompt_template_encode_end_idx ,:]
+
+            # Sample a random timestep for each image
+            grad_norm = None
+            with accelerator.accumulate(transformer):
+                # Predict the noise residual
+                optimizer.zero_grad()
+                # logit-normal
+                sigmas = torch.sigmoid(torch.randn((half_batch_size,), device=accelerator.device, dtype=latents.dtype))
+
+                if args.use_dynamic_shifting:
+                    sigmas = noise_scheduler.time_shift(mu, 1.0, sigmas)
+
+                sigmas = torch.concat( [sigmas,sigmas], dim=0 )
+                timesteps = sigmas * 1000.0
+                sigmas = sigmas.view(-1, 1, 1, 1)
+
+                noise = torch.randn_like(latents)
+                noise = noise.chunk(2)[0].repeat(2, 1, 1, 1)
+
+                noisy_latents = (1 - sigmas) * latents + sigmas * noise
+                noisy_latents = noisy_latents.to(weight_dtype)
+
+                packed_noisy_latents = pack_latents(
+                    noisy_latents,
+                    batch_size=latents.shape[0],
+                    num_channels_latents=latents.shape[1],
+                    height=latents.shape[2],
+                    width=latents.shape[3],
+                )
+
+                guidance = None
+                img_ids = prepare_pos_ids(modality_id=1,
+                                          type='image',
+                                          start=(prompt_embeds.shape[1], prompt_embeds.shape[1]),
+                                          height=latents.shape[2]//2,
+                                          width=latents.shape[3]//2).to(accelerator.device, dtype=torch.float64)
+
+                timesteps = (
+                    torch.tensor(timesteps)
+                    .expand(noisy_latents.shape[0])
+                    .to(device=accelerator.device)
+                    / 1000
+                )
+                text_ids = prepare_pos_ids(modality_id=0,
+                                           type='text',
+                                           start=(0, 0),
+                                           num_token=prompt_embeds.shape[1]).to(accelerator.device, torch.float64)
+
+                with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION):
+                    model_pred = transformer(packed_noisy_latents, prompt_embeds, timesteps,
+                                             img_ids, text_ids, guidance, return_dict=False)[0]
+
+                model_pred = unpack_latents(
+                    model_pred,
+                    height=latents.shape[2] * 8,
+                    width=latents.shape[3] * 8,
+                    vae_scale_factor=16,
+                )
+
+                target = noise - latents
+                model_losses = torch.mean(
+                    ((model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
+                    1,
+                )
+                model_loss_win, model_loss_lose = model_losses.chunk(2)
+                model_diff = model_loss_win - model_loss_lose
+                
+                with torch.no_grad():
+                    ref_pred = ref_model(packed_noisy_latents, prompt_embeds, timesteps,
+                                img_ids, text_ids, guidance, return_dict=False)[0]
+                
+                    ref_pred = unpack_latents(
+                        ref_pred,
+                        height=latents.shape[2] * 8,
+                        width=latents.shape[3] * 8,
+                        vae_scale_factor=16,
+                        )
+                    
+                    ref_losses = torch.mean(
+                            ((ref_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
+                            1,
+                        )
+                    ref_loss_win, ref_loss_lose = ref_losses.chunk(2)
+                    ref_diff = (ref_loss_win - ref_loss_lose) 
+
+                inside_term = -args.beta_dpo * (model_diff - ref_diff)
+                # implicit_acc = (inside_term > 0).sum().float() / inside_term.size(0)
+                loss = - F.logsigmoid(inside_term).mean()
+                
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    grad_norm = transformer.get_global_grad_norm()
+
+                optimizer.step()
+                if not accelerator.optimizer_step_was_skipped:  
+                    lr_scheduler.step()
+
+                if accelerator.sync_gradients and args.use_ema:
+                    model_ema.step(transformer.parameters())
+
+            lr = lr_scheduler.get_last_lr()[0]
+
+            if accelerator.sync_gradients:
+                global_step += 1
+
+                bsz, ic, ih, iw = image.shape
+                logs = {"loss": accelerator.gather(loss).mean().item(), 'aspect_ratio': (ih*1.0 / iw)}
+                if grad_norm is not None:
+                    logs.update(grad_norm=accelerator.gather(grad_norm).mean().item())
+
+                log_buffer.update(logs)
+                if (step + 1) % args.log_interval == 0 or (step + 1) == 1:
+                    t = (time.time() - last_tic) / args.log_interval
+                    t_d = data_time_all / args.log_interval
+
+                    log_buffer.average()
+                    info = f"Step={step+1}, Epoch={epoch}, global_step={global_step}, time_all:{t:.3f}, time_data:{t_d:.3f}, lr:{lr:.3e}, s:(ch:{latents.shape[1]},h:{latents.shape[2]},w:{latents.shape[3]}), "
+                    info += ', '.join([f"{k}:{v:.4f}" for k,v in log_buffer.output.items()])
+                    logger.info(info)
+                    last_tic = time.time()
+                    log_buffer.clear()
+                    data_time_all = 0
+                logs.update(lr=lr)
+                accelerator.log(logs, step=global_step)
+                data_time_start = time.time()
+
+                if global_step != 0 and global_step % args.save_model_steps == 0:
+                    save_path = os.path.join(
+                        args.work_dir, f'checkpoints-{global_step}')
+                    if args.use_ema:
+                        model_ema.store(transformer.parameters())
+                        model_ema.copy_to(transformer.parameters())
+
+                    accelerator.save_state(save_path)
+
+                    if args.use_ema:
+                        model_ema.restore(transformer.parameters())
+                    logger.info(f"Saved state to {save_path} (global_step: {global_step})")
+                    accelerator.wait_for_everyone()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Process some integers.")
+    parser.add_argument("--config", type=str, default='', help="config")
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        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(
+        "--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_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+    args = parse_args()
+
+    if args.config != '' and os.path.exists(args.config):
+        config = yaml.safe_load(open(args.config, 'r'))
+    else:
+        config = yaml.safe_load(open(f'{cur_dir}/train_config.yaml', 'r'))
+
+    args_dict = vars(args)
+    args_dict.update(config)
+    args = argparse.Namespace(**args_dict)
+
+    os.umask(0o000)
+    os.makedirs(args.work_dir, exist_ok=True)
+
+    log_dir = args.work_dir + f'/logs'
+    os.makedirs(log_dir, exist_ok=True)
+    accelerator_project_config = ProjectConfiguration(
+        project_dir=args.work_dir, logging_dir=log_dir)
+
+    with open(f'{log_dir}/train.yaml', 'w') as f:
+        yaml.dump(args_dict, f)
+
+    accelerator = Accelerator(
+        mixed_precision=args.mixed_precision,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        log_with=args.report_to,
+        project_config= accelerator_project_config,
+    )
+
+    # 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_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    logger.info(f'using weight_dtype {weight_dtype}!!!')
+
+    if args.diffusion_pretrain_weight:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(args.diffusion_pretrain_weight, ignore_mismatched_sizes=False)
+        ref_model = LongCatImageTransformer2DModel.from_pretrained(args.diffusion_pretrain_weight, ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight for BaseModel and RefModel, {args.diffusion_pretrain_weight}!!!')
+    else:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "transformer"), ignore_mismatched_sizes=False)
+        ref_model = LongCatImageTransformer2DModel.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "transformer"), ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight for BaseModel and RefModel, {args.pretrained_model_name_or_path+"/transformer"}!!!')
+
+    transformer = transformer.train()
+    ref_model.requires_grad_(False)
+    ref_model.to(accelerator.device, weight_dtype)
+    
+    total_trainable_params = sum(p.numel() for p in transformer.parameters() if p.requires_grad)
+    logger.info(f">>>>>> total_trainable_params: {total_trainable_params}")
+
+    if args.use_ema:
+        model_ema = EMAModel(transformer.parameters(), decay=args.ema_rate)
+    else:
+        model_ema = None
+
+    vae_dtype = torch.float32
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vae",  torch_dtype=weight_dtype).cuda().eval()
+
+    text_encoder = AutoModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder" , torch_dtype=weight_dtype, trust_remote_code=True).cuda().eval()
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="tokenizer" , torch_dtype=weight_dtype, trust_remote_code=True)
+    logger.info("all models loaded successfully")
+
+    # build models
+    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="scheduler")
+
+    latent_size = int(args.resolution) // 8
+    mu = calculate_shift(
+        (latent_size//2)**2,
+        noise_scheduler.config.base_image_seq_len,
+        noise_scheduler.config.max_image_seq_len,
+        noise_scheduler.config.base_shift,
+        noise_scheduler.config.max_shift,
+    )
+
+    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            for i, model in enumerate(models):
+                model.save_pretrained(os.path.join(output_dir, "transformer"))
+                if len(weights) != 0:
+                    weights.pop()
+
+    def load_model_hook(models, input_dir):
+        while len(models) > 0:
+            # pop models so that they are not loaded again
+            model = models.pop()
+            # load diffusers style into model
+            load_model = MeiGenImageTransformer2DModel.from_pretrained(
+                input_dir, subfolder="transformer")
+            model.register_to_config(**load_model.config)
+            model.load_state_dict(load_model.state_dict())
+            del load_model
+
+    accelerator.register_save_state_pre_hook(save_model_hook)
+    accelerator.register_load_state_pre_hook(load_model_hook)
+
+    if args.gradient_checkpointing:
+        transformer.enable_gradient_checkpointing()
+    
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
+        optimizer_class = bnb.optim.AdamW8bit
+    else:
+        optimizer_class = torch.optim.AdamW
+
+    params_to_optimize = transformer.parameters()
+    optimizer = optimizer_class(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    transformer.to(accelerator.device, dtype=weight_dtype)
+    if args.use_ema:
+        model_ema.to(accelerator.device, dtype=weight_dtype)
+
+    train_dataloader = build_dataloader(args, args.data_txt_root, tokenizer, args.resolution,)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+    )
+
+    global_step = 0
+    first_epoch = 0
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.work_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+        if path is None:
+            logger.info(f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.")
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            logger.info(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.work_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+    timestamp = time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())
+    if accelerator.is_main_process:
+        tracker_config = dict(vars(args))
+        try:
+            accelerator.init_trackers('dpo', tracker_config)
+        except Exception as e:
+            logger.warning(f'get error in save config, {e}')
+            accelerator.init_trackers(f"dpo_{timestamp}")
+
+    transformer, optimizer, _, _ = accelerator.prepare(
+        transformer, optimizer, train_dataloader, lr_scheduler)
+
+    log_buffer = LogBuffer()
+    train(global_step=global_step)

train_examples/edit/train.sh

@@ -0,0 +1,13 @@
+export TOKENIZERS_PARALLELISM=False
+export NCCL_DEBUG=INFO
+export NCCL_TIMEOUT=12000
+
+script_dir=$(cd -- "$(dirname -- "$0")" &> /dev/null && pwd -P)
+project_root=$(dirname "$(dirname "$script_dir")")
+echo "script_dir" ${script_dir}
+
+deepspeed_config_file=${project_root}/misc/accelerate_config.yaml
+
+accelerate launch  --mixed_precision bf16 --num_processes 8 --config_file ${deepspeed_config_file} \
+${script_dir}/train_edit.py \
+--config ${script_dir}/train_config.yaml

train_examples/edit/train_config.yaml

@@ -0,0 +1,46 @@
+# 1.Data setting
+data_txt_root: '/dataset/example/train_data_info.txt'  # data csv_filepath
+resolution: 1024 
+aspect_ratio_type: 'mar_1024'  # data bucketing strategy, mar_256、mar_512、mar_1024
+null_text_ratio: 0.1
+dataloader_num_workers: 8
+train_batch_size: 4
+repeats: 1
+
+prompt_template_encode_prefix: "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+prompt_template_encode_suffix: '<|im_end|>\n<|im_start|>assistant\n'
+prompt_template_encode_start_idx: 67
+prompt_template_encode_end_idx: 5
+
+# 2. Model setting
+text_tokenizer_max_length: 512  # tokenizer max len 
+pretrained_model_name_or_path: "./weights/LongCat-Image-Edit" # root directory of the model，with vae、transformer、scheduler eta;
+diffusion_pretrain_weight:  null # if a specified diffusion weight path is provided, load the model parameters from the current directory.
+use_dynamic_shifting: true #  scheduler dynamic shifting
+resume_from_checkpoint: latest
+  # - "latest"  # Loads most recent step checkpoint
+  # - "/path/to/checkpoint"  # Resumes from specified directory
+
+# 3. Training setting
+use_ema: False
+ema_rate: 0.999
+mixed_precision: 'bf16'
+max_train_steps: 100000
+gradient_accumulation_steps: 1
+gradient_checkpointing: true
+gradient_clip: 1.0
+learning_rate: 1.0e-5
+adam_weight_decay: 1.0e-2
+adam_epsilon: 1.0e-8
+adam_beta1: 0.9
+adam_beta2: 0.999
+lr_num_cycles: 1
+lr_power: 1.0
+lr_scheduler: 'constant'
+lr_warmup_steps: 1000
+
+#4. Log setting
+log_interval: 20
+save_model_steps: 1000
+work_dir: 'output/edit_model'
+seed: 43

train_examples/edit/train_dataset.py

@@ -0,0 +1,184 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+import os
+import random
+import traceback
+import math
+import json
+import numpy as np
+
+import torch
+import torch.distributed as dist
+import torchvision.transforms as T
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoTokenizer, AutoProcessor
+from PIL import Image
+from tqdm import tqdm
+
+from longcat_image.dataset import MULTI_RESOLUTION_MAP
+from longcat_image.utils import encode_prompt_edit
+from longcat_image.dataset import MultiResolutionDistributedSampler
+
+Image.MAX_IMAGE_PIXELS = 2000000000
+
+MAX_RETRY_NUMS = 100
+
+class Text2ImageLoraDataSet(torch.utils.data.Dataset):
+    def __init__(self,
+                 cfg: dict,
+                 txt_root: str,
+                 tokenizer: AutoTokenizer,
+                 text_processor:AutoProcessor,
+                 resolution: tuple = (1024, 1024),
+                 repeats: int = 1 ):
+        super(Text2ImageLoraDataSet, self).__init__()
+        self.resolution = resolution
+        self.text_tokenizer_max_length = cfg.text_tokenizer_max_length
+        self.null_text_ratio = cfg.null_text_ratio
+        self.aspect_ratio_type = cfg.aspect_ratio_type
+        self.aspect_ratio = MULTI_RESOLUTION_MAP[self.aspect_ratio_type]
+        self.tokenizer = tokenizer
+        self.image_processor_vl = text_processor.image_processor
+
+        self.prompt_template_encode_prefix = cfg.prompt_template_encode_prefix
+        self.prompt_template_encode_suffix = cfg.prompt_template_encode_suffix
+        self.prompt_template_encode_start_idx = cfg.prompt_template_encode_start_idx
+        self.prompt_template_encode_end_idx = cfg.prompt_template_encode_end_idx
+
+        self.total_datas = []
+        self.data_resolution_infos = []
+        with open(txt_root, 'r') as f:
+            lines = f.readlines()
+            lines *= cfg.repeats
+            for line in tqdm(lines):
+                data = json.loads(line.strip())
+                try:
+                    height, widht = int(data['height']), int(data['width'])
+                    self.data_resolution_infos.append((height, widht))
+                    self.total_datas.append(data)
+                except Exception as e:
+                    print(f'get error {e}, data {data}.')
+                    continue
+        self.data_nums = len(self.total_datas)
+        print(f'get sampler {len(self.total_datas)}, from {txt_root}!!!')
+
+    def transform_img(self, image, original_size, target_size):
+        img_h, img_w = original_size
+        target_height, target_width = target_size
+
+        original_aspect = img_h / img_w  # height/width
+        crop_aspect = target_height / target_width
+
+        if original_aspect >= crop_aspect:
+            resize_width = target_width
+            resize_height = math.ceil(img_h * (target_width/img_w))
+        else:
+            resize_width = math.ceil(img_w * (target_height/img_h))
+            resize_height = target_height
+
+        image = T.Compose([
+            T.Resize((resize_height, resize_width),interpolation=InterpolationMode.BICUBIC), 
+            T.CenterCrop((target_height, target_width)),
+            T.ToTensor(),
+            T.Normalize([.5], [.5]),
+        ])(image)
+
+        return image
+    def transform_img_vl(self, image, original_size, target_size):
+        img_h, img_w = original_size
+        target_height, target_width = target_size
+
+        original_aspect = img_h / img_w  # height/width
+        crop_aspect = target_height / target_width
+
+        if original_aspect >= crop_aspect:
+            resize_width = target_width
+            resize_height = math.ceil(img_h * (target_width/img_w))
+        else:
+            resize_width = math.ceil(img_w * (target_height/img_h))
+            resize_height = target_height
+
+        image = T.Compose([
+            T.Resize((resize_height, resize_width),interpolation=InterpolationMode.BICUBIC), 
+            T.CenterCrop((target_height, target_width)),
+            T.Resize((target_height//2, target_width//2)),
+        ])(image)
+
+        return image
+
+    def __getitem__(self, index_tuple):
+        index, target_size = index_tuple
+
+        for _ in range(MAX_RETRY_NUMS):
+            try:
+                item = self.total_datas[index]
+                img_path = item["img_path"]
+                ref_img_path = item["ref_img_path"]
+                prompt = item['prompt']
+
+                if random.random() < self.null_text_ratio:
+                    prompt = ''
+
+                raw_image = Image.open(img_path).convert('RGB')
+                ref_image = Image.open(ref_img_path).convert('RGB')
+                assert raw_image is not None
+                img_w, img_h = raw_image.size
+
+                ref_image_vl = self.transform_img_vl(ref_image, original_size=(img_h, img_w), target_size= target_size )
+                raw_image = self.transform_img(raw_image, original_size=(img_h, img_w), target_size= target_size )
+                ref_image = self.transform_img(ref_image, original_size=(img_h, img_w), target_size= target_size )
+                
+                input_ids, attention_mask, pixel_values, image_grid_thw = encode_prompt_edit(prompt,ref_image_vl, self.tokenizer, self.image_processor_vl,self.text_tokenizer_max_length, self.prompt_template_encode_prefix, self.prompt_template_encode_suffix )
+                return {"image": raw_image, "ref_image":ref_image, "prompt": prompt, 'input_ids': input_ids, 'attention_mask': attention_mask, 'pixel_values':pixel_values, 'image_grid_thw':image_grid_thw}
+
+            except Exception as e:
+                traceback.print_exc()
+                print(f"failed read data {e}!!!")
+                index = random.randint(0, self.data_nums-1)
+
+    def __len__(self):
+        return self.data_nums
+
+    def collate_fn(self, batchs):
+        images = torch.stack([example["image"] for example in batchs])
+        ref_images = torch.stack([example["ref_image"] for example in batchs])
+        input_ids = torch.stack([example["input_ids"] for example in batchs])
+        attention_mask = torch.stack([example["attention_mask"] for example in batchs])
+        pixel_values = torch.stack([example["pixel_values"] for example in batchs])
+        image_grid_thw = torch.stack([example["image_grid_thw"] for example in batchs])
+        prompts = [example['prompt'] for example in batchs]
+        batch_dict = {
+            "images": images,
+            "ref_images": ref_images,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "prompts": prompts,
+            "pixel_values":pixel_values,
+            "image_grid_thw":image_grid_thw
+        }
+        return batch_dict
+
+
+def build_dataloader(cfg: dict,
+                     csv_root: str,
+                     tokenizer: AutoTokenizer,
+                     text_processor: AutoProcessor,
+                     resolution: tuple = (1024, 1024)):
+    dataset = Text2ImageLoraDataSet(cfg, csv_root, tokenizer, text_processor, resolution)
+
+    sampler = MultiResolutionDistributedSampler(batch_size=cfg.train_batch_size, dataset=dataset,
+                                                data_resolution_infos=dataset.data_resolution_infos,
+                                                bucket_info=dataset.aspect_ratio,
+                                                epoch=0,
+                                                num_replicas=None,
+                                                rank=None
+                                                )
+
+    train_loader = torch.utils.data.DataLoader(
+        dataset,
+        collate_fn=dataset.collate_fn,
+        batch_size=cfg.train_batch_size,
+        num_workers=cfg.dataloader_num_workers,
+        sampler=sampler,
+        shuffle=None,
+    )
+    return train_loader

train_examples/edit/train_edit.py

@@ -0,0 +1,460 @@
+import os
+import sys
+import time
+import warnings
+
+import argparse
+import yaml
+import torch
+import math
+import logging
+import transformers
+import diffusers
+from pathlib import Path
+from transformers import Qwen2Model, Qwen2TokenizerFast
+from accelerate import Accelerator, InitProcessGroupKwargs
+from accelerate.utils import ProjectConfiguration, set_seed
+from accelerate.logging import get_logger
+from diffusers.models import AutoencoderKL
+from diffusers.optimization import get_scheduler
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.training_utils import EMAModel
+from diffusers.utils.import_utils import is_xformers_available
+from transformers import AutoTokenizer, AutoModel, AutoProcessor
+
+from train_dataset import build_dataloader
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.utils import LogBuffer
+from longcat_image.utils import pack_latents, unpack_latents, calculate_shift, prepare_pos_ids
+
+
+warnings.filterwarnings("ignore")  # ignore warning
+
+current_file_path = Path(__file__).resolve()
+sys.path.insert(0, str(current_file_path.parent.parent))
+
+logger = get_logger(__name__)
+
+def train(global_step=0):
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    
+    # Train!
+    total_batch_size = args.train_batch_size * \
+        accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataloader.dataset)}")
+    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    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}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+
+    last_tic = time.time()
+
+    # Now you train the model
+    for epoch in range(first_epoch, args.num_train_epochs + 1):
+        data_time_start = time.time()
+        data_time_all = 0
+
+        for step, batch in enumerate(train_dataloader):
+            image = batch['images']
+            ref_image = batch['ref_images']
+
+            data_time_all += time.time() - data_time_start
+
+            with torch.no_grad():
+                latents = vae.encode(image.to(weight_dtype).to(accelerator.device)).latent_dist.sample()
+                latents = latents.to(dtype=(weight_dtype))
+                latents = (latents - vae.config.shift_factor) * vae.config.scaling_factor
+
+                ref_latents = vae.encode(ref_image.to(weight_dtype).to(accelerator.device)).latent_dist.sample()
+                ref_latents = ref_latents.to(dtype=(weight_dtype))
+                ref_latents = (ref_latents - vae.config.shift_factor) * vae.config.scaling_factor
+
+            text_input_ids = batch['input_ids'].to(accelerator.device)
+            text_attention_mask = batch['attention_mask'].to(accelerator.device)
+            pixel_values = batch['pixel_values'].to(accelerator.device)
+            image_grid_thw = batch['image_grid_thw'].to(accelerator.device)
+
+            with torch.no_grad():
+                text_output = text_encoder(
+                    input_ids=text_input_ids,
+                    attention_mask=text_attention_mask,
+                    pixel_values=pixel_values,
+                    image_grid_thw=image_grid_thw,
+                    output_hidden_states=True
+                )
+                prompt_embeds = text_output.hidden_states[-1].clone().detach()
+
+            prompt_embeds = prompt_embeds.to(weight_dtype)
+            prompt_embeds = prompt_embeds[:,args.prompt_template_encode_start_idx: -args.prompt_template_encode_end_idx ,:]
+
+            # Sample a random timestep for each image
+            grad_norm = None
+            with accelerator.accumulate(transformer):
+                # Predict the noise residual
+                optimizer.zero_grad()
+                # logit-normal
+                sigmas = torch.sigmoid(torch.randn((latents.shape[0],), device=accelerator.device, dtype=latents.dtype))
+
+                if args.use_dynamic_shifting:
+                    sigmas = noise_scheduler.time_shift(mu, 1.0, sigmas)
+
+                timesteps = sigmas * 1000.0
+                sigmas = sigmas.view(-1, 1, 1, 1)
+
+                noise = torch.randn_like(latents)
+
+                noisy_latents = (1 - sigmas) * latents + sigmas * noise
+                noisy_latents = noisy_latents.to(weight_dtype)
+
+                packed_noisy_latents = pack_latents(
+                    noisy_latents,
+                    batch_size=latents.shape[0],
+                    num_channels_latents=latents.shape[1],
+                    height=latents.shape[2],
+                    width=latents.shape[3],
+                )
+                
+                packed_ref_latents = pack_latents(
+                    ref_latents,
+                    batch_size=ref_latents.shape[0],
+                    num_channels_latents=ref_latents.shape[1],
+                    height=ref_latents.shape[2],
+                    width=ref_latents.shape[3],
+                )
+
+                guidance = None
+                img_ids = prepare_pos_ids(modality_id=1,
+                                          type='image',
+                                          start=(prompt_embeds.shape[1], prompt_embeds.shape[1]),
+                                          height=latents.shape[2]//2,
+                                          width=latents.shape[3]//2).to(accelerator.device, dtype=torch.float64)
+                img_ids_ref = prepare_pos_ids(modality_id=2,
+                                          type='image',
+                                          start=(prompt_embeds.shape[1], prompt_embeds.shape[1]),
+                                          height=ref_latents.shape[2]//2,
+                                          width=ref_latents.shape[3]//2).to(accelerator.device, dtype=torch.float64)
+
+                timesteps = (
+                    torch.tensor(timesteps)
+                    .expand(noisy_latents.shape[0])
+                    .to(device=accelerator.device)
+                    / 1000
+                )
+                text_ids = prepare_pos_ids(modality_id=0,
+                                           type='text',
+                                           start=(0, 0),
+                                           num_token=prompt_embeds.shape[1]).to(accelerator.device, torch.float64)
+
+                img_ids = torch.cat([img_ids, img_ids_ref], dim=0)
+                latent_model_input = torch.cat([packed_noisy_latents, packed_ref_latents], dim=1)
+                with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION):
+                    model_pred = transformer(latent_model_input, prompt_embeds, timesteps,
+                                             img_ids, text_ids, guidance, return_dict=False)[0]
+                    model_pred = model_pred[:, :packed_noisy_latents.size(1)]
+
+                model_pred = unpack_latents(
+                    model_pred,
+                    height=latents.shape[2] * 8,
+                    width=latents.shape[3] * 8,
+                    vae_scale_factor=16,
+                )
+
+                target = noise - latents
+                loss = torch.mean(
+                    ((model_pred.float() - target.float()) ** 2).reshape(
+                        target.shape[0], -1
+                    ),
+                    1,
+                ).mean()
+
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    grad_norm = transformer.get_global_grad_norm()
+
+                optimizer.step()
+                if not accelerator.optimizer_step_was_skipped:  
+                    lr_scheduler.step()
+
+                if accelerator.sync_gradients and args.use_ema:
+                    model_ema.step(transformer.parameters())
+
+            lr = lr_scheduler.get_last_lr()[0]
+
+            if accelerator.sync_gradients:
+                bsz, ic, ih, iw = image.shape
+                logs = {"loss": accelerator.gather(loss).mean().item(), 'aspect_ratio': (ih*1.0 / iw)}
+                if grad_norm is not None:
+                    logs.update(grad_norm=accelerator.gather(grad_norm).mean().item())
+
+                log_buffer.update(logs)
+                if (step + 1) % args.log_interval == 0 or (step + 1) == 1:
+                    t = (time.time() - last_tic) / args.log_interval
+                    t_d = data_time_all / args.log_interval
+
+                    log_buffer.average()
+                    info = f"Step={step+1}, Epoch={epoch}, global_step={global_step}, time_all:{t:.3f}, time_data:{t_d:.3f}, lr:{lr:.3e}, s:(ch:{latents.shape[1]},h:{latents.shape[2]},w:{latents.shape[3]}), "
+                    info += ', '.join([f"{k}:{v:.4f}" for k,v in log_buffer.output.items()])
+                    logger.info(info)
+                    last_tic = time.time()
+                    log_buffer.clear()
+                    data_time_all = 0
+                logs.update(lr=lr)
+                accelerator.log(logs, step=global_step)
+                global_step += 1
+                data_time_start = time.time()
+
+                if global_step != 0 and global_step % args.save_model_steps == 0:
+                    save_path = os.path.join(args.work_dir, f'checkpoints-{global_step}')
+                    if args.use_ema:
+                        model_ema.store(transformer.parameters())
+                        model_ema.copy_to(transformer.parameters())
+
+                    accelerator.save_state(save_path)
+
+                    if args.use_ema:
+                        model_ema.restore(transformer.parameters())
+                    logger.info(f"Saved state to {save_path} (global_step: {global_step})")
+                    accelerator.wait_for_everyone()
+
+            if global_step >= args.max_train_steps:
+                break
+        
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Process some integers.")
+    parser.add_argument("--config", type=str, default='', help="config")
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        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(
+        "--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_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+    args = parse_args()
+
+    if args.config != '' and os.path.exists(args.config):
+        config = yaml.safe_load(open(args.config, 'r'))
+    else:
+        config = yaml.safe_load(open(f'{cur_dir}/train_config.yaml', 'r'))
+
+    args_dict = vars(args)
+    args_dict.update(config)
+    args = argparse.Namespace(**args_dict)
+
+    os.umask(0o000)
+    os.makedirs(args.work_dir, exist_ok=True)
+
+    log_dir = args.work_dir + f'/logs'
+    os.makedirs(log_dir, exist_ok=True)
+    accelerator_project_config = ProjectConfiguration(project_dir=args.work_dir, logging_dir=log_dir)
+
+    with open(f'{log_dir}/train.yaml', 'w') as f:
+        yaml.dump(args_dict, f)
+
+    accelerator = Accelerator(
+        mixed_precision=args.mixed_precision,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        log_with=args.report_to,
+        project_config= accelerator_project_config,
+    )
+
+    # 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_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    logger.info(f'using weight_dtype {weight_dtype}!!!')
+
+    if args.diffusion_pretrain_weight:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(args.diffusion_pretrain_weight, ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight {args.diffusion_pretrain_weight}!!!')
+    else:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "transformer"), ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight {args.pretrained_model_name_or_path+"/transformer"}!!!')
+
+    transformer = transformer.train()
+
+    total_trainable_params = sum(p.numel() for p in transformer.parameters() if p.requires_grad)
+    logger.info(f">>>>>> total_trainable_params: {total_trainable_params}")
+
+    if args.use_ema:
+        model_ema = EMAModel(transformer.parameters(), decay=args.ema_rate)
+    else:
+        model_ema = None
+
+    vae_dtype = torch.float32
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vae",  torch_dtype=weight_dtype).cuda().eval()
+
+    text_encoder = AutoModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder" , torch_dtype=weight_dtype, trust_remote_code=True).cuda().eval()
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="tokenizer" , torch_dtype=weight_dtype, trust_remote_code=True)
+    text_processor = AutoProcessor.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="tokenizer" , torch_dtype=weight_dtype, trust_remote_code=True)
+    logger.info("all models loaded successfully")
+
+    # build models
+    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="scheduler")
+
+    latent_size = int(args.resolution) // 8
+    mu = calculate_shift(
+        (latent_size//2)**2,
+        noise_scheduler.config.base_image_seq_len,
+        noise_scheduler.config.max_image_seq_len,
+        noise_scheduler.config.base_shift,
+        noise_scheduler.config.max_shift,
+    )
+
+    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            for i, model in enumerate(models):
+                model.save_pretrained(os.path.join(output_dir, "transformer"))
+                if len(weights) != 0:
+                    weights.pop()
+
+    def load_model_hook(models, input_dir):
+        while len(models) > 0:
+            # pop models so that they are not loaded again
+            model = models.pop()
+            # load diffusers style into model
+            load_model = LongCatImageTransformer2DModel.from_pretrained(
+                input_dir, subfolder="transformer")
+            model.register_to_config(**load_model.config)
+            model.load_state_dict(load_model.state_dict())
+            del load_model
+
+    accelerator.register_save_state_pre_hook(save_model_hook)
+    accelerator.register_load_state_pre_hook(load_model_hook)
+
+    if args.gradient_checkpointing:
+        transformer.enable_gradient_checkpointing()
+
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
+        optimizer_class = bnb.optim.AdamW8bit
+    else:
+        optimizer_class = torch.optim.AdamW
+
+    params_to_optimize = transformer.parameters()
+    optimizer = optimizer_class(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    transformer.to(accelerator.device, dtype=weight_dtype)
+    if args.use_ema:
+        model_ema.to(accelerator.device, dtype=weight_dtype)
+
+    train_dataloader = build_dataloader(args, args.data_txt_root, tokenizer, text_processor,args.resolution,)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+    )
+
+    global_step = 0
+    first_epoch = 0
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.work_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+        if path is None:
+            logger.info(f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.")
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            logger.info(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.work_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+    timestamp = time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())
+    if accelerator.is_main_process:
+        tracker_config = dict(vars(args))
+        try:
+            accelerator.init_trackers('sft', tracker_config)
+        except Exception as e:
+            logger.warning(f'get error in save config, {e}')
+            accelerator.init_trackers(f"sft_{timestamp}")
+
+    transformer, optimizer, _, _ = accelerator.prepare(
+        transformer, optimizer, train_dataloader, lr_scheduler)
+
+    log_buffer = LogBuffer()
+    train(global_step=global_step)

train_examples/lora/train.sh

@@ -0,0 +1,13 @@
+export TOKENIZERS_PARALLELISM=False
+export NCCL_DEBUG=INFO
+export NCCL_TIMEOUT=12000
+
+script_dir=$(cd -- "$(dirname -- "$0")" &> /dev/null && pwd -P)
+project_root=$(dirname "$(dirname "$script_dir")")
+echo "script_dir" ${script_dir}
+
+deepspeed_config_file=${project_root}/misc/accelerate_config.yaml
+
+accelerate launch  --mixed_precision bf16 --num_processes 8 --config_file ${deepspeed_config_file} \
+${script_dir}/train_lora.py \
+--config ${script_dir}/train_config.yaml

train_examples/lora/train_config.yaml

@@ -0,0 +1,48 @@
+# 1.Data setting
+
+data_txt_root: '/dataset/example/train_data_info.txt'  # data csv_filepath
+resolution: 1024 
+aspect_ratio_type: 'mar_1024'  # data bucketing strategy, mar_256、mar_512、mar_1024
+null_text_ratio: 0.1
+dataloader_num_workers: 8
+train_batch_size: 2
+repeats: 100
+
+prompt_template_encode_prefix: '<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n'
+prompt_template_encode_suffix: '<|im_end|>\n<|im_start|>assistant\n'
+prompt_template_encode_start_idx: 36
+prompt_template_encode_end_idx: 5
+
+# 2. Model setting
+text_tokenizer_max_length: 512  # tokenizer max len 
+pretrained_model_name_or_path: "/xxx/weights/Longcat-Image-Dev" # root directory of the model，with vae、transformer、scheduler eta;
+diffusion_pretrain_weight:  null # if a specified diffusion weight path is provided, load the model parameters from the current directory.
+use_dynamic_shifting: true #  scheduler dynamic shifting
+resume_from_checkpoint: latest
+  # - "latest"  # Loads most recent step checkpoint
+  # - "/path/to/checkpoint"  # Resumes from specified directory
+
+# 3. Training setting
+lora_rank: 32
+use_ema: False
+ema_rate: 0.999
+mixed_precision: 'bf16'
+max_train_steps: 100000
+gradient_accumulation_steps: 1
+gradient_checkpointing: true
+gradient_clip: 1.0
+learning_rate: 1.0e-4
+adam_weight_decay: 1.0e-3
+adam_epsilon: 1.0e-8
+adam_beta1: 0.9
+adam_beta2: 0.999
+lr_num_cycles: 1
+lr_power: 1.0
+lr_scheduler: 'constant'
+lr_warmup_steps: 1000
+
+#4. Log setting
+log_interval: 10
+save_model_steps: 200
+work_dir: 'output/lora_model'
+seed: 43

train_examples/lora/train_dataset.py

@@ -0,0 +1,197 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+import os
+import random
+import traceback
+import math
+import json
+import numpy as np
+
+import torch
+import torch.distributed as dist
+import torchvision.transforms as T
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoTokenizer
+from PIL import Image
+from tqdm import tqdm
+
+from longcat_image.dataset import MULTI_RESOLUTION_MAP
+from longcat_image.utils import encode_prompt
+from longcat_image.dataset import MultiResolutionDistributedSampler
+
+Image.MAX_IMAGE_PIXELS = 2000000000
+
+MAX_RETRY_NUMS = 100
+
+class Text2ImageLoraDataSet(torch.utils.data.Dataset):
+    def __init__(self,
+                 cfg: dict,
+                 txt_root: str,
+                 tokenizer: AutoTokenizer,
+                 resolution: tuple = (1024, 1024),
+                 repeats: int = 1 ):
+        super(Text2ImageLoraDataSet, self).__init__()
+        self.resolution = resolution
+        self.text_tokenizer_max_length = cfg.text_tokenizer_max_length
+        self.null_text_ratio = cfg.null_text_ratio
+        self.aspect_ratio_type = cfg.aspect_ratio_type
+        self.aspect_ratio = MULTI_RESOLUTION_MAP[self.aspect_ratio_type]
+        self.tokenizer = tokenizer
+
+        self.prompt_template_encode_prefix = cfg.prompt_template_encode_prefix
+        self.prompt_template_encode_suffix = cfg.prompt_template_encode_suffix
+        self.prompt_template_encode_start_idx = cfg.prompt_template_encode_start_idx
+        self.prompt_template_encode_end_idx = cfg.prompt_template_encode_end_idx
+
+        self.total_datas = []
+        self.data_resolution_infos = []
+        with open(txt_root, 'r') as f:
+            lines = f.readlines()
+            lines *= cfg.repeats
+            for line in tqdm(lines):
+                data = json.loads(line.strip())
+                try:
+                    height, widht = int(data['height']), int(data['width'])
+                    self.data_resolution_infos.append((height, widht))
+                    self.total_datas.append(data)
+                except Exception as e:
+                    print(f'get error {e}, data {data}.')
+                    continue
+        self.data_nums = len(self.total_datas)
+        print(f'get sampler {len(self.total_datas)}, from {txt_root}!!!')
+
+    def transform_img(self, image, original_size, target_size):
+        img_h, img_w = original_size
+        target_height, target_width = target_size
+
+        original_aspect = img_h / img_w  # height/width
+        crop_aspect = target_height / target_width
+
+        if original_aspect >= crop_aspect:
+            resize_width = target_width
+            resize_height = math.ceil(img_h * (target_width/img_w))
+        else:
+            resize_width = math.ceil(img_w * (target_height/img_h))
+            resize_height = target_height
+
+        image = T.Compose([
+            T.Resize((resize_height, resize_width),interpolation=InterpolationMode.BICUBIC),  # Image.LANCZOS
+            T.CenterCrop((target_height, target_width)),
+            T.ToTensor(),
+            T.Normalize([.5], [.5]),
+        ])(image)
+
+        return image
+
+    def __getitem__(self, index_tuple):
+        index, target_size = index_tuple
+
+        for _ in range(MAX_RETRY_NUMS):
+            try:
+                item = self.total_datas[index]
+                img_path = item["img_path"]
+                prompt = item['prompt']
+
+                if random.random() < self.null_text_ratio:
+                    prompt = ''
+
+                raw_image = Image.open(img_path).convert('RGB')
+                assert raw_image is not None
+                img_w, img_h = raw_image.size
+
+                raw_image = self.transform_img(raw_image, original_size=(img_h, img_w), target_size= target_size )
+                input_ids,attention_mask = encode_prompt(prompt, self.tokenizer, self.text_tokenizer_max_length, self.prompt_template_encode_prefix, self.prompt_template_encode_suffix )
+                return {"image": raw_image, "prompt": prompt, 'input_ids': input_ids, 'attention_mask': attention_mask}
+
+            except Exception as e:
+                traceback.print_exc()
+                print(f"failed read data {e}!!!")
+                index = random.randint(0, self.data_nums-1)
+
+    def __len__(self):
+        return self.data_nums
+
+    def collate_fn(self, batchs):
+        images = torch.stack([example["image"] for example in batchs])
+        input_ids = torch.stack([example["input_ids"] for example in batchs])
+        attention_mask = torch.stack([example["attention_mask"] for example in batchs])
+        prompts = [example['prompt'] for example in batchs]
+        batch_dict = {
+            "images": images,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "prompts": prompts,
+        }
+        return batch_dict
+
+
+def build_dataloader(cfg: dict,
+                     csv_root: str,
+                     tokenizer: AutoTokenizer,
+                     resolution: tuple = (1024, 1024)):
+    dataset = Text2ImageLoraDataSet(cfg, csv_root, tokenizer, resolution)
+
+    sampler = MultiResolutionDistributedSampler(batch_size=cfg.train_batch_size, dataset=dataset,
+                                                data_resolution_infos=dataset.data_resolution_infos,
+                                                bucket_info=dataset.aspect_ratio,
+                                                epoch=0,
+                                                num_replicas=None,
+                                                rank=None
+                                                )
+
+    train_loader = torch.utils.data.DataLoader(
+        dataset,
+        collate_fn=dataset.collate_fn,
+        batch_size=cfg.train_batch_size,
+        num_workers=cfg.dataloader_num_workers,
+        sampler=sampler,
+        shuffle=None,
+    )
+    return train_loader
+
+
+if __name__ == '__main__':
+    import sys
+    import argparse
+    from torchvision.transforms.functional import to_pil_image
+
+    txt_root = 'xxx'
+    cfg = argparse.Namespace(
+        txt_root=txt_root,
+        text_tokenizer_max_length=512,
+        resolution=1024,
+        text_encoder_path="xxx",
+        center_crop=True,
+        dataloader_num_workers=0,
+        null_text_ratio=0.1,
+        train_batch_size=16,
+        seed=0,
+        aspect_ratio_type='mar_1024',
+        revision=None)
+
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(cfg.text_encoder_path, trust_remote_code=True)
+    data_loader = build_dataloader(cfg, cfg.csv_root, tokenizer, cfg.resolution)
+
+    _oroot = f'./debug_data_example_show'
+    os.makedirs(_oroot, exist_ok=True)
+
+    cnt = 0
+    for epoch in range(1):
+        print(f"Start, epoch {epoch}!!!")
+        for i_batch, batch in enumerate(data_loader):
+            print(batch['attention_mask'].shape)
+            print(batch['images'].shape)
+
+            batch_prompts = batch['prompts']
+            for idx, per_img in enumerate(batch['images']):
+                re_transforms = T.Compose([
+                    T.Normalize(mean=[-0.5/0.5], std=[1.0/0.5])
+                ])
+                prompt = batch_prompts[idx]
+                img = to_pil_image(re_transforms(per_img))
+                prompt = prompt[:min(30, len(prompt))]
+                oname = _oroot + f'/{str(i_batch)}_{str(idx)}_{prompt}.png'
+                img.save(oname)
+            if cnt > 100:
+                break
+            cnt += 1

train_examples/lora/train_lora.py

@@ -0,0 +1,461 @@
+import os
+import sys
+import time
+import warnings
+
+import argparse
+import yaml
+import torch
+import math
+import logging
+import transformers
+import diffusers
+from pathlib import Path
+from transformers import Qwen2Model, Qwen2TokenizerFast
+from accelerate import Accelerator, InitProcessGroupKwargs
+from accelerate.utils import ProjectConfiguration, set_seed
+from accelerate.logging import get_logger
+from diffusers.models import AutoencoderKL
+from diffusers.optimization import get_scheduler
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.training_utils import EMAModel
+from diffusers.utils.import_utils import is_xformers_available
+from transformers import AutoTokenizer, AutoModel
+
+from train_dataset import build_dataloader
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.utils import LogBuffer
+from longcat_image.utils import pack_latents, unpack_latents, calculate_shift, prepare_pos_ids
+from peft import LoraConfig, set_peft_model_state_dict
+from peft import LoraConfig, get_peft_model
+from peft.utils import get_peft_model_state_dict
+
+warnings.filterwarnings("ignore")  # ignore warning
+
+current_file_path = Path(__file__).resolve()
+sys.path.insert(0, str(current_file_path.parent.parent))
+
+logger = get_logger(__name__)
+
+def train(global_step=0):
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    
+    # Train!
+    total_batch_size = args.train_batch_size * \
+        accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataloader.dataset)}")
+    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    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}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+
+    last_tic = time.time()
+
+    # Now you train the model
+    for epoch in range(first_epoch, args.num_train_epochs + 1):
+        data_time_start = time.time()
+        data_time_all = 0
+
+        for step, batch in enumerate(train_dataloader):
+            image = batch['images']
+
+            data_time_all += time.time() - data_time_start
+
+            with torch.no_grad():
+                latents = vae.encode(image.to(weight_dtype).to(accelerator.device)).latent_dist.sample()
+                latents = latents.to(dtype=(weight_dtype))
+                latents = (latents - vae.config.shift_factor) * vae.config.scaling_factor
+
+            text_input_ids = batch['input_ids'].to(accelerator.device)
+            text_attention_mask = batch['attention_mask'].to(accelerator.device)
+
+            with torch.no_grad():
+                text_output = text_encoder(
+                    input_ids=text_input_ids,
+                    attention_mask=text_attention_mask,
+                    output_hidden_states=True
+                )
+                prompt_embeds = text_output.hidden_states[-1].clone().detach()
+
+            prompt_embeds = prompt_embeds.to(weight_dtype)
+            prompt_embeds = prompt_embeds[:,args.prompt_template_encode_start_idx: -args.prompt_template_encode_end_idx ,:]
+
+            # Sample a random timestep for each image
+            grad_norm = None
+            with accelerator.accumulate(transformer):
+                # Predict the noise residual
+                optimizer.zero_grad()
+                # logit-normal
+                sigmas = torch.sigmoid(torch.randn((latents.shape[0],), device=accelerator.device, dtype=latents.dtype))
+
+                if args.use_dynamic_shifting:
+                    sigmas = noise_scheduler.time_shift(mu, 1.0, sigmas)
+
+                timesteps = sigmas * 1000.0
+                sigmas = sigmas.view(-1, 1, 1, 1)
+
+                noise = torch.randn_like(latents)
+
+                noisy_latents = (1 - sigmas) * latents + sigmas * noise
+                noisy_latents = noisy_latents.to(weight_dtype)
+
+                packed_noisy_latents = pack_latents(
+                    noisy_latents,
+                    batch_size=latents.shape[0],
+                    num_channels_latents=latents.shape[1],
+                    height=latents.shape[2],
+                    width=latents.shape[3],
+                )
+
+                guidance = None
+                img_ids = prepare_pos_ids(modality_id=1,
+                                          type='image',
+                                          start=(prompt_embeds.shape[1], prompt_embeds.shape[1]),
+                                          height=latents.shape[2]//2,
+                                          width=latents.shape[3]//2).to(accelerator.device, dtype=torch.float64)
+
+                timesteps = (
+                    torch.tensor(timesteps)
+                    .expand(noisy_latents.shape[0])
+                    .to(device=accelerator.device)
+                    / 1000
+                )
+                text_ids = prepare_pos_ids(modality_id=0,
+                                           type='text',
+                                           start=(0, 0),
+                                           num_token=prompt_embeds.shape[1]).to(accelerator.device, torch.float64)
+
+                with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION):
+                    model_pred = transformer(packed_noisy_latents, prompt_embeds, timesteps,
+                                             img_ids, text_ids, guidance, return_dict=False)[0]
+
+                model_pred = unpack_latents(
+                    model_pred,
+                    height=latents.shape[2] * 8,
+                    width=latents.shape[3] * 8,
+                    vae_scale_factor=16,
+                )
+
+                target = noise - latents
+                loss = torch.mean(
+                    ((model_pred.float() - target.float()) ** 2).reshape(
+                        target.shape[0], -1
+                    ),
+                    1,
+                ).mean()
+
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    grad_norm = transformer.get_global_grad_norm()
+
+                optimizer.step()
+                if not accelerator.optimizer_step_was_skipped:  
+                    lr_scheduler.step()
+
+                if accelerator.sync_gradients and args.use_ema:
+                    model_ema.step(transformer.parameters())
+
+            lr = lr_scheduler.get_last_lr()[0]
+
+            if accelerator.sync_gradients:
+                bsz, ic, ih, iw = image.shape
+                logs = {"loss": accelerator.gather(loss).mean().item(), 'aspect_ratio': (ih*1.0 / iw)}
+                if grad_norm is not None:
+                    logs.update(grad_norm=accelerator.gather(grad_norm).mean().item())
+
+                log_buffer.update(logs)
+                if (step + 1) % args.log_interval == 0 or (step + 1) == 1:
+                    t = (time.time() - last_tic) / args.log_interval
+                    t_d = data_time_all / args.log_interval
+
+                    log_buffer.average()
+                    info = f"Step={step+1}, Epoch={epoch}, global_step={global_step}, time_all:{t:.3f}, time_data:{t_d:.3f}, lr:{lr:.3e}, s:(ch:{latents.shape[1]},h:{latents.shape[2]},w:{latents.shape[3]}), "
+                    info += ', '.join([f"{k}:{v:.4f}" for k,v in log_buffer.output.items()])
+                    logger.info(info)
+                    last_tic = time.time()
+                    log_buffer.clear()
+                    data_time_all = 0
+                logs.update(lr=lr)
+                accelerator.log(logs, step=global_step)
+                global_step += 1
+                data_time_start = time.time()
+
+                if global_step != 0 and global_step % args.save_model_steps == 0:
+                    accelerator.wait_for_everyone()
+                    os.umask(0o000)
+                    cur_lora_ckpt_save_dir = f"{args.work_dir}/checkpoints-{global_step}"
+                    os.makedirs(cur_lora_ckpt_save_dir, exist_ok=True)
+                    if accelerator.is_main_process:
+                        if hasattr(transformer, 'module'):
+                            transformer.module.save_pretrained(cur_lora_ckpt_save_dir)
+                        else:
+                            transformer.save_pretrained(cur_lora_ckpt_save_dir)
+                        logger.info(f'Saved lora checkpoint of epoch {epoch} to {cur_lora_ckpt_save_dir}.')
+                    accelerator.save_state(cur_lora_ckpt_save_dir)
+                    accelerator.wait_for_everyone()
+
+            if global_step >= args.max_train_steps:
+                break
+        
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Process some integers.")
+    parser.add_argument("--config", type=str, default='', help="config")
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        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(
+        "--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_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+    args = parse_args()
+
+    if args.config != '' and os.path.exists(args.config):
+        config = yaml.safe_load(open(args.config, 'r'))
+    else:
+        config = yaml.safe_load(open(f'{cur_dir}/train_config.yaml', 'r'))
+
+    args_dict = vars(args)
+    args_dict.update(config)
+    args = argparse.Namespace(**args_dict)
+
+    os.umask(0o000)
+    os.makedirs(args.work_dir, exist_ok=True)
+
+    log_dir = args.work_dir + f'/logs'
+    os.makedirs(log_dir, exist_ok=True)
+    accelerator_project_config = ProjectConfiguration(project_dir=args.work_dir, logging_dir=log_dir)
+
+    with open(f'{log_dir}/train.yaml', 'w') as f:
+        yaml.dump(args_dict, f)
+
+    accelerator = Accelerator(
+        mixed_precision=args.mixed_precision,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        log_with=args.report_to,
+        project_config= accelerator_project_config,
+    )
+
+    # 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_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    logger.info(f'using weight_dtype {weight_dtype}!!!')
+
+    if args.diffusion_pretrain_weight:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(args.diffusion_pretrain_weight, ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight {args.diffusion_pretrain_weight}!!!')
+    else:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "transformer"), ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight {args.pretrained_model_name_or_path+"/transformer"}!!!')
+
+    # transformer = transformer.train()
+    
+    target_modules = [
+            "attn.to_k",
+            "attn.to_q",
+            "attn.to_v",
+            "attn.to_out.0",
+            "attn.add_k_proj",
+            "attn.add_q_proj",
+            "attn.add_v_proj",
+            "attn.to_add_out",
+            "ff.net.0.proj",
+            "ff.net.2",
+            "ff_context.net.0.proj",
+            "ff_context.net.2",
+        ]
+
+    lora_config = LoraConfig(
+        r=args.lora_rank,
+        init_lora_weights="gaussian",
+        target_modules= target_modules ,
+        use_dora=False,
+        use_rslora=False
+    )
+    transformer = get_peft_model(transformer, lora_config)
+    transformer.print_trainable_parameters()
+
+    total_trainable_params = sum(p.numel() for p in transformer.parameters() if p.requires_grad)
+    logger.info(f">>>>>> total_trainable_params: {total_trainable_params}")
+    
+    if args.use_ema:
+        model_ema = EMAModel(transformer.parameters(), decay=args.ema_rate)
+    else:
+        model_ema = None
+
+    vae_dtype = torch.float32
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vae",  torch_dtype=weight_dtype).cuda().eval()
+
+    text_encoder = AutoModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder" , torch_dtype=weight_dtype, trust_remote_code=True).cuda().eval()
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="tokenizer" , torch_dtype=weight_dtype, trust_remote_code=True)
+    logger.info("all models loaded successfully")
+
+    # build models
+    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="scheduler")
+
+    latent_size = int(args.resolution) // 8
+    mu = calculate_shift(
+        (latent_size//2)**2,
+        noise_scheduler.config.base_image_seq_len,
+        noise_scheduler.config.max_image_seq_len,
+        noise_scheduler.config.base_shift,
+        noise_scheduler.config.max_shift,
+    )
+
+    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            for i, model in enumerate(models):
+                model.save_pretrained(os.path.join(output_dir, "transformer"))
+                if len(weights) != 0:
+                    weights.pop()
+
+    def load_model_hook(models, input_dir):
+        while len(models) > 0:
+            # pop models so that they are not loaded again
+            model = models.pop()
+            # load diffusers style into model
+            load_model = LongCatImageTransformer2DModel.from_pretrained(
+                input_dir, subfolder="transformer")
+            model.register_to_config(**load_model.config)
+            model.load_state_dict(load_model.state_dict())
+            del load_model
+
+    accelerator.register_save_state_pre_hook(save_model_hook)
+    accelerator.register_load_state_pre_hook(load_model_hook)
+
+    if args.gradient_checkpointing:
+        transformer.enable_gradient_checkpointing()
+
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
+        optimizer_class = bnb.optim.AdamW8bit
+    else:
+        optimizer_class = torch.optim.AdamW
+
+    params_to_optimize = transformer.parameters()
+    optimizer = optimizer_class(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    transformer.to(accelerator.device, dtype=weight_dtype)
+    if args.use_ema:
+        model_ema.to(accelerator.device, dtype=weight_dtype)
+
+    train_dataloader = build_dataloader(args, args.data_txt_root, tokenizer, args.resolution,)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+    )
+
+    global_step = 0
+    first_epoch = 0
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.work_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+        if path is None:
+            logger.info(f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.")
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            logger.info(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.work_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+    timestamp = time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())
+    if accelerator.is_main_process:
+        tracker_config = dict(vars(args))
+        try:
+            accelerator.init_trackers('sft', tracker_config)
+        except Exception as e:
+            logger.warning(f'get error in save config, {e}')
+            accelerator.init_trackers(f"sft_{timestamp}")
+
+    transformer, optimizer, _, _ = accelerator.prepare(
+        transformer, optimizer, train_dataloader, lr_scheduler)
+
+    log_buffer = LogBuffer()
+    train(global_step=global_step)

train_examples/sft/train.sh

@@ -0,0 +1,13 @@
+export TOKENIZERS_PARALLELISM=False
+export NCCL_DEBUG=INFO
+export NCCL_TIMEOUT=12000
+
+script_dir=$(cd -- "$(dirname -- "$0")" &> /dev/null && pwd -P)
+project_root=$(dirname "$(dirname "$script_dir")")
+echo "script_dir" ${script_dir}
+
+deepspeed_config_file=${project_root}/misc/accelerate_config.yaml
+
+accelerate launch  --mixed_precision bf16 --num_processes 8 --config_file ${deepspeed_config_file} \
+${script_dir}/train_sft.py \
+--config ${script_dir}/train_config.yaml

train_examples/sft/train_config.yaml

@@ -0,0 +1,46 @@
+# 1.Data setting
+data_txt_root: '/dataset/example/train_data_info.txt'  # data csv_filepath
+resolution: 1024 
+aspect_ratio_type: 'mar_1024'  # data bucketing strategy, mar_256、mar_512、mar_1024
+null_text_ratio: 0.1
+dataloader_num_workers: 8
+train_batch_size: 4
+repeats: 1
+
+prompt_template_encode_prefix: '<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n'
+prompt_template_encode_suffix: '<|im_end|>\n<|im_start|>assistant\n'
+prompt_template_encode_start_idx: 36
+prompt_template_encode_end_idx: 5
+
+# 2. Model setting
+text_tokenizer_max_length: 512  # tokenizer max len 
+pretrained_model_name_or_path: "/xxx/weights/Longcat-Image-Dev" # root directory of the model，with vae、transformer、scheduler eta;
+diffusion_pretrain_weight:  null # if a specified diffusion weight path is provided, load the model parameters from the current directory.
+use_dynamic_shifting: true #  scheduler dynamic shifting
+resume_from_checkpoint: latest
+  # - "latest"  # Loads most recent step checkpoint
+  # - "/path/to/checkpoint"  # Resumes from specified directory
+
+# 3. Training setting
+use_ema: False
+ema_rate: 0.999
+mixed_precision: 'bf16'
+max_train_steps: 100000
+gradient_accumulation_steps: 1
+gradient_checkpointing: true
+gradient_clip: 1.0
+learning_rate: 1.0e-5
+adam_weight_decay: 1.0e-2
+adam_epsilon: 1.0e-8
+adam_beta1: 0.9
+adam_beta2: 0.999
+lr_num_cycles: 1
+lr_power: 1.0
+lr_scheduler: 'constant'
+lr_warmup_steps: 1000
+
+#4. Log setting
+log_interval: 20
+save_model_steps: 1000
+work_dir: 'output/sft_model'
+seed: 43

train_examples/sft/train_dataset.py

@@ -0,0 +1,197 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+import os
+import random
+import traceback
+import math
+import json
+import numpy as np
+
+import torch
+import torch.distributed as dist
+import torchvision.transforms as T
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoTokenizer
+from PIL import Image
+from tqdm import tqdm
+
+from longcat_image.dataset import MULTI_RESOLUTION_MAP
+from longcat_image.utils import encode_prompt
+from longcat_image.dataset import MultiResolutionDistributedSampler
+
+Image.MAX_IMAGE_PIXELS = 2000000000
+
+MAX_RETRY_NUMS = 100
+
+class Text2ImageLoraDataSet(torch.utils.data.Dataset):
+    def __init__(self,
+                 cfg: dict,
+                 txt_root: str,
+                 tokenizer: AutoTokenizer,
+                 resolution: tuple = (1024, 1024),
+                 repeats: int = 1 ):
+        super(Text2ImageLoraDataSet, self).__init__()
+        self.resolution = resolution
+        self.text_tokenizer_max_length = cfg.text_tokenizer_max_length
+        self.null_text_ratio = cfg.null_text_ratio
+        self.aspect_ratio_type = cfg.aspect_ratio_type
+        self.aspect_ratio = MULTI_RESOLUTION_MAP[self.aspect_ratio_type]
+        self.tokenizer = tokenizer
+
+        self.prompt_template_encode_prefix = cfg.prompt_template_encode_prefix
+        self.prompt_template_encode_suffix = cfg.prompt_template_encode_suffix
+        self.prompt_template_encode_start_idx = cfg.prompt_template_encode_start_idx
+        self.prompt_template_encode_end_idx = cfg.prompt_template_encode_end_idx
+
+        self.total_datas = []
+        self.data_resolution_infos = []
+        with open(txt_root, 'r') as f:
+            lines = f.readlines()
+            lines *= cfg.repeats
+            for line in tqdm(lines):
+                data = json.loads(line.strip())
+                try:
+                    height, widht = int(data['height']), int(data['width'])
+                    self.data_resolution_infos.append((height, widht))
+                    self.total_datas.append(data)
+                except Exception as e:
+                    print(f'get error {e}, data {data}.')
+                    continue
+        self.data_nums = len(self.total_datas)
+        print(f'get sampler {len(self.total_datas)}, from {txt_root}!!!')
+
+    def transform_img(self, image, original_size, target_size):
+        img_h, img_w = original_size
+        target_height, target_width = target_size
+
+        original_aspect = img_h / img_w  # height/width
+        crop_aspect = target_height / target_width
+
+        if original_aspect >= crop_aspect:
+            resize_width = target_width
+            resize_height = math.ceil(img_h * (target_width/img_w))
+        else:
+            resize_width = math.ceil(img_w * (target_height/img_h))
+            resize_height = target_height
+
+        image = T.Compose([
+            T.Resize((resize_height, resize_width),interpolation=InterpolationMode.BICUBIC),  # Image.LANCZOS
+            T.CenterCrop((target_height, target_width)),
+            T.ToTensor(),
+            T.Normalize([.5], [.5]),
+        ])(image)
+
+        return image
+
+    def __getitem__(self, index_tuple):
+        index, target_size = index_tuple
+
+        for _ in range(MAX_RETRY_NUMS):
+            try:
+                item = self.total_datas[index]
+                img_path = item["img_path"]
+                prompt = item['prompt']
+
+                if random.random() < self.null_text_ratio:
+                    prompt = ''
+
+                raw_image = Image.open(img_path).convert('RGB')
+                assert raw_image is not None
+                img_w, img_h = raw_image.size
+
+                raw_image = self.transform_img(raw_image, original_size=(img_h, img_w), target_size= target_size )
+                input_ids,attention_mask = encode_prompt(prompt, self.tokenizer, self.text_tokenizer_max_length, self.prompt_template_encode_prefix, self.prompt_template_encode_suffix )
+                return {"image": raw_image, "prompt": prompt, 'input_ids': input_ids, 'attention_mask': attention_mask}
+
+            except Exception as e:
+                traceback.print_exc()
+                print(f"failed read data {e}!!!")
+                index = random.randint(0, self.data_nums-1)
+
+    def __len__(self):
+        return self.data_nums
+
+    def collate_fn(self, batchs):
+        images = torch.stack([example["image"] for example in batchs])
+        input_ids = torch.stack([example["input_ids"] for example in batchs])
+        attention_mask = torch.stack([example["attention_mask"] for example in batchs])
+        prompts = [example['prompt'] for example in batchs]
+        batch_dict = {
+            "images": images,
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "prompts": prompts,
+        }
+        return batch_dict
+
+
+def build_dataloader(cfg: dict,
+                     csv_root: str,
+                     tokenizer: AutoTokenizer,
+                     resolution: tuple = (1024, 1024)):
+    dataset = Text2ImageLoraDataSet(cfg, csv_root, tokenizer, resolution)
+
+    sampler = MultiResolutionDistributedSampler(batch_size=cfg.train_batch_size, dataset=dataset,
+                                                data_resolution_infos=dataset.data_resolution_infos,
+                                                bucket_info=dataset.aspect_ratio,
+                                                epoch=0,
+                                                num_replicas=None,
+                                                rank=None
+                                                )
+
+    train_loader = torch.utils.data.DataLoader(
+        dataset,
+        collate_fn=dataset.collate_fn,
+        batch_size=cfg.train_batch_size,
+        num_workers=cfg.dataloader_num_workers,
+        sampler=sampler,
+        shuffle=None,
+    )
+    return train_loader
+
+
+if __name__ == '__main__':
+    import sys
+    import argparse
+    from torchvision.transforms.functional import to_pil_image
+
+    txt_root = 'xxx'
+    cfg = argparse.Namespace(
+        txt_root=txt_root,
+        text_tokenizer_max_length=512,
+        resolution=1024,
+        text_encoder_path="xxx",
+        center_crop=True,
+        dataloader_num_workers=0,
+        null_text_ratio=0.1,
+        train_batch_size=16,
+        seed=0,
+        aspect_ratio_type='mar_1024',
+        revision=None)
+
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(cfg.text_encoder_path, trust_remote_code=True)
+    data_loader = build_dataloader(cfg, cfg.csv_root, tokenizer, cfg.resolution)
+
+    _oroot = f'./debug_data_example_show'
+    os.makedirs(_oroot, exist_ok=True)
+
+    cnt = 0
+    for epoch in range(1):
+        print(f"Start, epoch {epoch}!!!")
+        for i_batch, batch in enumerate(data_loader):
+            print(batch['attention_mask'].shape)
+            print(batch['images'].shape)
+
+            batch_prompts = batch['prompts']
+            for idx, per_img in enumerate(batch['images']):
+                re_transforms = T.Compose([
+                    T.Normalize(mean=[-0.5/0.5], std=[1.0/0.5])
+                ])
+                prompt = batch_prompts[idx]
+                img = to_pil_image(re_transforms(per_img))
+                prompt = prompt[:min(30, len(prompt))]
+                oname = _oroot + f'/{str(i_batch)}_{str(idx)}_{prompt}.png'
+                img.save(oname)
+            if cnt > 100:
+                break
+            cnt += 1

train_examples/sft/train_sft.py

@@ -0,0 +1,433 @@
+import os
+import sys
+import time
+import warnings
+
+import argparse
+import yaml
+import torch
+import math
+import logging
+import transformers
+import diffusers
+from pathlib import Path
+from transformers import Qwen2Model, Qwen2TokenizerFast
+from accelerate import Accelerator, InitProcessGroupKwargs
+from accelerate.utils import ProjectConfiguration, set_seed
+from accelerate.logging import get_logger
+from diffusers.models import AutoencoderKL
+from diffusers.optimization import get_scheduler
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.training_utils import EMAModel
+from diffusers.utils.import_utils import is_xformers_available
+from transformers import AutoTokenizer, AutoModel
+
+from train_dataset import build_dataloader
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.utils import LogBuffer
+from longcat_image.utils import pack_latents, unpack_latents, calculate_shift, prepare_pos_ids
+
+
+warnings.filterwarnings("ignore")  # ignore warning
+
+current_file_path = Path(__file__).resolve()
+sys.path.insert(0, str(current_file_path.parent.parent))
+
+logger = get_logger(__name__)
+
+def train(global_step=0):
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    
+    # Train!
+    total_batch_size = args.train_batch_size * \
+        accelerator.num_processes * args.gradient_accumulation_steps
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataloader.dataset)}")
+    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    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}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+
+    last_tic = time.time()
+
+    # Now you train the model
+    for epoch in range(first_epoch, args.num_train_epochs + 1):
+        data_time_start = time.time()
+        data_time_all = 0
+
+        for step, batch in enumerate(train_dataloader):
+            image = batch['images']
+
+            data_time_all += time.time() - data_time_start
+
+            with torch.no_grad():
+                latents = vae.encode(image.to(weight_dtype).to(accelerator.device)).latent_dist.sample()
+                latents = latents.to(dtype=(weight_dtype))
+                latents = (latents - vae.config.shift_factor) * vae.config.scaling_factor
+
+            text_input_ids = batch['input_ids'].to(accelerator.device)
+            text_attention_mask = batch['attention_mask'].to(accelerator.device)
+
+            with torch.no_grad():
+                text_output = text_encoder(
+                    input_ids=text_input_ids,
+                    attention_mask=text_attention_mask,
+                    output_hidden_states=True
+                )
+                prompt_embeds = text_output.hidden_states[-1].clone().detach()
+
+            prompt_embeds = prompt_embeds.to(weight_dtype)
+            prompt_embeds = prompt_embeds[:,args.prompt_template_encode_start_idx: -args.prompt_template_encode_end_idx ,:]
+
+            # Sample a random timestep for each image
+            grad_norm = None
+            with accelerator.accumulate(transformer):
+                # Predict the noise residual
+                optimizer.zero_grad()
+                # logit-normal
+                sigmas = torch.sigmoid(torch.randn((latents.shape[0],), device=accelerator.device, dtype=latents.dtype))
+
+                if args.use_dynamic_shifting:
+                    sigmas = noise_scheduler.time_shift(mu, 1.0, sigmas)
+
+                timesteps = sigmas * 1000.0
+                sigmas = sigmas.view(-1, 1, 1, 1)
+
+                noise = torch.randn_like(latents)
+
+                noisy_latents = (1 - sigmas) * latents + sigmas * noise
+                noisy_latents = noisy_latents.to(weight_dtype)
+
+                packed_noisy_latents = pack_latents(
+                    noisy_latents,
+                    batch_size=latents.shape[0],
+                    num_channels_latents=latents.shape[1],
+                    height=latents.shape[2],
+                    width=latents.shape[3],
+                )
+
+                guidance = None
+                img_ids = prepare_pos_ids(modality_id=1,
+                                          type='image',
+                                          start=(prompt_embeds.shape[1], prompt_embeds.shape[1]),
+                                          height=latents.shape[2]//2,
+                                          width=latents.shape[3]//2).to(accelerator.device, dtype=torch.float64)
+
+                timesteps = (
+                    torch.tensor(timesteps)
+                    .expand(noisy_latents.shape[0])
+                    .to(device=accelerator.device)
+                    / 1000
+                )
+                text_ids = prepare_pos_ids(modality_id=0,
+                                           type='text',
+                                           start=(0, 0),
+                                           num_token=prompt_embeds.shape[1]).to(accelerator.device, torch.float64)
+
+                with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION):
+                    model_pred = transformer(packed_noisy_latents, prompt_embeds, timesteps,
+                                             img_ids, text_ids, guidance, return_dict=False)[0]
+
+                model_pred = unpack_latents(
+                    model_pred,
+                    height=latents.shape[2] * 8,
+                    width=latents.shape[3] * 8,
+                    vae_scale_factor=16,
+                )
+
+                target = noise - latents
+                loss = torch.mean(
+                    ((model_pred.float() - target.float()) ** 2).reshape(
+                        target.shape[0], -1
+                    ),
+                    1,
+                ).mean()
+
+                accelerator.backward(loss)
+
+                if accelerator.sync_gradients:
+                    grad_norm = transformer.get_global_grad_norm()
+
+                optimizer.step()
+                if not accelerator.optimizer_step_was_skipped:  
+                    lr_scheduler.step()
+
+                if accelerator.sync_gradients and args.use_ema:
+                    model_ema.step(transformer.parameters())
+
+            lr = lr_scheduler.get_last_lr()[0]
+
+            if accelerator.sync_gradients:
+                bsz, ic, ih, iw = image.shape
+                logs = {"loss": accelerator.gather(loss).mean().item(), 'aspect_ratio': (ih*1.0 / iw)}
+                if grad_norm is not None:
+                    logs.update(grad_norm=accelerator.gather(grad_norm).mean().item())
+
+                log_buffer.update(logs)
+                if (step + 1) % args.log_interval == 0 or (step + 1) == 1:
+                    t = (time.time() - last_tic) / args.log_interval
+                    t_d = data_time_all / args.log_interval
+
+                    log_buffer.average()
+                    info = f"Step={step+1}, Epoch={epoch}, global_step={global_step}, time_all:{t:.3f}, time_data:{t_d:.3f}, lr:{lr:.3e}, s:(ch:{latents.shape[1]},h:{latents.shape[2]},w:{latents.shape[3]}), "
+                    info += ', '.join([f"{k}:{v:.4f}" for k,v in log_buffer.output.items()])
+                    logger.info(info)
+                    last_tic = time.time()
+                    log_buffer.clear()
+                    data_time_all = 0
+                logs.update(lr=lr)
+                accelerator.log(logs, step=global_step)
+                global_step += 1
+                data_time_start = time.time()
+
+                if global_step != 0 and global_step % args.save_model_steps == 0:
+                    save_path = os.path.join(args.work_dir, f'checkpoints-{global_step}')
+                    if args.use_ema:
+                        model_ema.store(transformer.parameters())
+                        model_ema.copy_to(transformer.parameters())
+
+                    accelerator.save_state(save_path)
+
+                    if args.use_ema:
+                        model_ema.restore(transformer.parameters())
+                    logger.info(f"Saved state to {save_path} (global_step: {global_step})")
+                    accelerator.wait_for_everyone()
+
+            if global_step >= args.max_train_steps:
+                break
+        
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Process some integers.")
+    parser.add_argument("--config", type=str, default='', help="config")
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        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(
+        "--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_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+    args = parse_args()
+
+    if args.config != '' and os.path.exists(args.config):
+        config = yaml.safe_load(open(args.config, 'r'))
+    else:
+        config = yaml.safe_load(open(f'{cur_dir}/train_config.yaml', 'r'))
+
+    args_dict = vars(args)
+    args_dict.update(config)
+    args = argparse.Namespace(**args_dict)
+
+    os.umask(0o000)
+    os.makedirs(args.work_dir, exist_ok=True)
+
+    log_dir = args.work_dir + f'/logs'
+    os.makedirs(log_dir, exist_ok=True)
+    accelerator_project_config = ProjectConfiguration(project_dir=args.work_dir, logging_dir=log_dir)
+
+    with open(f'{log_dir}/train.yaml', 'w') as f:
+        yaml.dump(args_dict, f)
+
+    accelerator = Accelerator(
+        mixed_precision=args.mixed_precision,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        log_with=args.report_to,
+        project_config= accelerator_project_config,
+    )
+
+    # 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_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    logger.info(f'using weight_dtype {weight_dtype}!!!')
+
+    if args.diffusion_pretrain_weight:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(args.diffusion_pretrain_weight, ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight {args.diffusion_pretrain_weight}!!!')
+    else:
+        transformer = LongCatImageTransformer2DModel.from_pretrained(os.path.join(args.pretrained_model_name_or_path, "transformer"), ignore_mismatched_sizes=False)
+        logger.info(f'successful load model weight {args.pretrained_model_name_or_path+"/transformer"}!!!')
+
+    transformer = transformer.train()
+
+    total_trainable_params = sum(p.numel() for p in transformer.parameters() if p.requires_grad)
+    logger.info(f">>>>>> total_trainable_params: {total_trainable_params}")
+
+    if args.use_ema:
+        model_ema = EMAModel(transformer.parameters(), decay=args.ema_rate)
+    else:
+        model_ema = None
+
+    vae_dtype = torch.float32
+    vae = AutoencoderKL.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vae",  torch_dtype=weight_dtype).cuda().eval()
+
+    text_encoder = AutoModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder" , torch_dtype=weight_dtype, trust_remote_code=True).cuda().eval()
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="tokenizer" , torch_dtype=weight_dtype, trust_remote_code=True)
+    logger.info("all models loaded successfully")
+
+    # build models
+    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="scheduler")
+
+    latent_size = int(args.resolution) // 8
+    mu = calculate_shift(
+        (latent_size//2)**2,
+        noise_scheduler.config.base_image_seq_len,
+        noise_scheduler.config.max_image_seq_len,
+        noise_scheduler.config.base_shift,
+        noise_scheduler.config.max_shift,
+    )
+
+    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            for i, model in enumerate(models):
+                model.save_pretrained(os.path.join(output_dir, "transformer"))
+                if len(weights) != 0:
+                    weights.pop()
+
+    def load_model_hook(models, input_dir):
+        while len(models) > 0:
+            # pop models so that they are not loaded again
+            model = models.pop()
+            # load diffusers style into model
+            load_model = LongCatImageTransformer2DModel.from_pretrained(
+                input_dir, subfolder="transformer")
+            model.register_to_config(**load_model.config)
+            model.load_state_dict(load_model.state_dict())
+            del load_model
+
+    accelerator.register_save_state_pre_hook(save_model_hook)
+    accelerator.register_load_state_pre_hook(load_model_hook)
+
+    if args.gradient_checkpointing:
+        transformer.enable_gradient_checkpointing()
+
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
+        optimizer_class = bnb.optim.AdamW8bit
+    else:
+        optimizer_class = torch.optim.AdamW
+
+    params_to_optimize = transformer.parameters()
+    optimizer = optimizer_class(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    transformer.to(accelerator.device, dtype=weight_dtype)
+    if args.use_ema:
+        model_ema.to(accelerator.device, dtype=weight_dtype)
+
+    train_dataloader = build_dataloader(args, args.data_txt_root, tokenizer, args.resolution,)
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+    )
+
+    global_step = 0
+    first_epoch = 0
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.work_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+        if path is None:
+            logger.info(f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.")
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            logger.info(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.work_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+    timestamp = time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())
+    if accelerator.is_main_process:
+        tracker_config = dict(vars(args))
+        try:
+            accelerator.init_trackers('sft', tracker_config)
+        except Exception as e:
+            logger.warning(f'get error in save config, {e}')
+            accelerator.init_trackers(f"sft_{timestamp}")
+
+    transformer, optimizer, _, _ = accelerator.prepare(
+        transformer, optimizer, train_dataloader, lr_scheduler)
+
+    log_buffer = LogBuffer()
+    train(global_step=global_step)