Add project files

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)