train/train_text_decoder.py

# Copyright 2024 The HuggingFace Team and The MeissonFlow Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import copy
import logging
import math
import os
import json
from contextlib import nullcontext
from pathlib import Path
import sys
sys.path.append(os.getcwd())
import gc
import torch
import torch.nn.functional as F
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from peft import LoraConfig
from peft.utils import get_peft_model_state_dict
from torch.utils.data import DataLoader, default_collate
from torchvision import transforms
from transformers import (
    CLIPTextModelWithProjection,
    CLIPTokenizer,
)
import diffusers.optimization
from diffusers import EMAModel, VQModel 
from src.scheduler import Scheduler
from diffusers.loaders import LoraLoaderMixin
from diffusers.utils import is_wandb_available
from src.pipeline import UnifiedPipeline
from torchvision.utils import save_image,make_grid
from datasets import load_dataset
from train.trainer_utils import save_checkpoint
from train.dataset_utils import ImageCaptionDataset, HuggingFaceDataset
from train.dataset_utils import tokenize_prompt, encode_prompt
from src.transformer import SymmetricTransformer2DModel, Transformer2DModel

if is_wandb_available():
    import wandb
    # wandb.login(key="")

logger = get_logger(__name__, log_level="INFO")

import torch._dynamo
torch._dynamo.config.verbose = True

# Optionally suppress errors to fall back to eager execution
torch._dynamo.config.suppress_errors = True

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--pretrained_model_architecture",
        type=str,
        default="Meissonic",
        required=False
    )
    parser.add_argument(
        "--text_encoder_architecture",
        type=str,
        default="open_clip",
        required=False,
        help="The architecture of the text encoder. One of ['CLIP', 'open_clip', 'flan-t5-base','Qwen2-0.5B','gemini-2b',long_CLIP_T5_base','CLIP_T5_base']",
    )
    parser.add_argument(
        "--instance_dataset",
        type=str,
        default=None,
        required=False,
        help="The dataset to use for training. One of ['MSCOCO600K', 'PickaPicV2']",
    )
    parser.add_argument(
        "--training_from_scratch",
        type=bool,
        default=False,
        required=False
    )
    parser.add_argument(
        "--pretrained_model_name_or_path",
        type=str,
        default=None,
        required=True,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--revision",
        type=str,
        default=None,
        required=False,
        help="Revision of pretrained model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--variant",
        type=str,
        default=None,
        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
    )
    parser.add_argument(
        "--instance_data_dataset",
        type=str,
        default=None,
        required=False,
        help="A Hugging Face dataset containing the training images",
    )
    parser.add_argument(
        "--instance_data_dir",
        type=str,
        default=None,
        required=False,
        help="A folder containing the training data of instance images.",
    )
    parser.add_argument(
        "--instance_data_image", type=str, default=None, required=False, help="A single training image"
    )
    parser.add_argument(
        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
    )
    parser.add_argument(
        "--dataloader_num_workers",
        type=int,
        default=0,
        help=(
            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
        ),
    )
    parser.add_argument(
        "--allow_tf32",
        action="store_true",
        help=(
            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
        ),
    )
    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
    parser.add_argument("--ema_decay", type=float, default=0.9999)
    parser.add_argument("--ema_update_after_step", type=int, default=0)
    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument(
        "--output_dir",
        type=str,
        default="muse_training",
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
    parser.add_argument(
        "--logging_dir",
        type=str,
        default="logs",
        help=(
            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
        ),
    )
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--checkpointing_steps",
        type=int,
        default=500,
        help=(
            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
            "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."
            "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."
            "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"
            "instructions."
        ),
    )
    parser.add_argument(
        "--logging_steps",
        type=int,
        default=50,
    )
    parser.add_argument(
        "--checkpoints_total_limit",
        type=int,
        default=None,
        help=(
            "Max number of checkpoints to store. Passed as `total_limit` to the `Accelerator` `ProjectConfiguration`."
            " See Accelerator::save_state https://huggingface.co/docs/accelerate/package_reference/accelerator#accelerate.Accelerator.save_state"
            " for more details"
        ),
    )
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help=(
            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
        ),
    )
    parser.add_argument(
        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=0.0003,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--scale_lr",
        action="store_true",
        default=False,
        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
    )
    parser.add_argument(
        "--lr_scheduler",
        type=str,
        default="constant",
        help=(
            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
            ' "constant", "constant_with_warmup"]'
        ),
    )
    parser.add_argument(
        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument(
        "--validation_steps",
        type=int,
        default=100,
        help=(
            "Run validation every X steps. Validation consists of running the prompt"
            " `args.validation_prompt` multiple times: `args.num_validation_images`"
            " and logging the images."
        ),
    )
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default=None,
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
        ),
    )
    parser.add_argument(
        "--report_to",
        type=str,
        default="wandb",
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
        ),
    )
    parser.add_argument("--validation_prompts", type=str, nargs="*")
    parser.add_argument(
        "--resolution",
        type=int,
        default=512,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
    parser.add_argument("--split_vae_encode", type=int, required=False, default=None)
    parser.add_argument("--min_masking_rate", type=float, default=0.0)
    parser.add_argument("--cond_dropout_prob", type=float, default=0.0)
    parser.add_argument("--max_grad_norm", default=50.0, type=float, help="Max gradient norm.", required=False)
    parser.add_argument("--use_lora", action="store_true", help="Fine tune the model using LoRa")
    parser.add_argument("--text_encoder_use_lora", action="store_true", help="Fine tune the model using LoRa")
    parser.add_argument("--lora_r", default=16, type=int)
    parser.add_argument("--lora_alpha", default=32, type=int)
    parser.add_argument("--lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
    parser.add_argument("--text_encoder_lora_r", default=16, type=int)
    parser.add_argument("--text_encoder_lora_alpha", default=32, type=int)
    parser.add_argument("--text_encoder_lora_target_modules", default=["to_q", "to_k", "to_v"], type=str, nargs="+")
    parser.add_argument("--train_text_encoder", action="store_true")
    parser.add_argument("--image_key", type=str, required=False)
    parser.add_argument("--prompt_key", type=str, required=False)
    parser.add_argument(
        "--gradient_checkpointing",
        action="store_true",
        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
    )
    parser.add_argument("--prompt_prefix", type=str, required=False, default=None)

    args = parser.parse_args()

    if args.report_to == "wandb":
        if not is_wandb_available():
            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")

    num_datasources = sum(
        [x is not None for x in [args.instance_data_dir, args.instance_data_image, args.instance_data_dataset]]
    )

    if num_datasources != 1:
        raise ValueError(
            "provide one and only one of `--instance_data_dir`, `--instance_data_image`, or `--instance_data_dataset`"
        )

    if args.instance_data_dir is not None:
        if not os.path.exists(args.instance_data_dir):
            raise ValueError(f"Does not exist: `--args.instance_data_dir` {args.instance_data_dir}")

    if args.instance_data_image is not None:
        if not os.path.exists(args.instance_data_image):
            raise ValueError(f"Does not exist: `--args.instance_data_image` {args.instance_data_image}")

    if args.instance_data_dataset is not None and (args.image_key is None or args.prompt_key is None):
        raise ValueError("`--instance_data_dataset` requires setting `--image_key` and `--prompt_key`")

    return args

def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
    latent_image_ids = torch.zeros(height // 2, width // 2, 3)
    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None]
    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :]

    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape

    latent_image_ids = latent_image_ids.reshape(
        latent_image_id_height * latent_image_id_width, latent_image_id_channels
    )
    # latent_image_ids = latent_image_ids.unsqueeze(0).repeat(batch_size, 1, 1)

    return latent_image_ids.to(device=device, dtype=dtype)

def main(args):
    if args.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True

    # if args.pretrained_model_architecture == "Meissonic":
    #     from src.pipeline import Pipeline
    # else:
    #     raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")


    logging_dir = Path(args.output_dir, args.logging_dir)

    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)

    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.report_to,
        project_config=accelerator_project_config,
    )

    if accelerator.is_main_process:
        os.makedirs(args.output_dir, exist_ok=True)

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)

    if accelerator.is_main_process:
        accelerator.init_trackers("meissonic", config=vars(copy.deepcopy(args)))

    if args.seed is not None:
        set_seed(args.seed)

    if args.text_encoder_architecture == "open_clip":
        if args.resume_from_checkpoint:
            text_encoder = CLIPTextModelWithProjection.from_pretrained( 
                args.resume_from_checkpoint, subfolder="text_encoder", variant=args.variant
            )
            tokenizer = CLIPTokenizer.from_pretrained(
                args.resume_from_checkpoint, subfolder="tokenizer", variant=args.variant
            )
        else:
            text_encoder = CLIPTextModelWithProjection.from_pretrained(
                args.pretrained_model_name_or_path, subfolder="text_encoder", variant=args.variant
            )
            tokenizer = CLIPTokenizer.from_pretrained(
                args.pretrained_model_name_or_path, subfolder="tokenizer", variant=args.variant
            )
    else:
        raise ValueError(f"Unknown text encoder architecture: {args.text_encoder_architecture}")
    
    vq_model = VQModel.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vqvae", revision=args.revision, variant=args.variant
    )

    if args.train_text_encoder:
        if args.text_encoder_use_lora:
            lora_config = LoraConfig(
                r=args.text_encoder_lora_r,
                lora_alpha=args.text_encoder_lora_alpha,
                target_modules=args.text_encoder_lora_target_modules,
            )
            text_encoder.add_adapter(lora_config)
        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
            text_encoder[0].train()
            text_encoder[0].requires_grad_(True)
            text_encoder[1].train()
            text_encoder[1].requires_grad_(True)
        else:
            text_encoder.train()
            text_encoder.requires_grad_(True)
    else:
        if args.text_encoder_architecture == "CLIP_T5_base":  # Not support yet. Only support open_clip
            text_encoder[0].eval()
            text_encoder[0].requires_grad_(False)
            text_encoder[1].eval()
            text_encoder[1].requires_grad_(False)
        else:
            text_encoder.eval()
            text_encoder.requires_grad_(False)

    vq_model.requires_grad_(False)

    if args.pretrained_model_architecture == "Meissonic":
        if args.training_from_scratch:
            model = SymmetricTransformer2DModel( 
                patch_size = 1,
                in_channels = 64,
                num_layers = 14, 
                num_single_layers = 28,
                attention_head_dim = 128, 
                num_attention_heads = 8, 
                joint_attention_dim = 1024,
                pooled_projection_dim = 1024,
                guidance_embeds = False,
                axes_dims_rope = (16, 56, 56),
                downsample= True,
                upsample= True,
            )
        else:
            orig_model = Transformer2DModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="transformer", low_cpu_mem_usage=False, device_map=None)
            orig_config = orig_model.config
            config = {}
            for k, v in orig_config.items():
                if k.startswith("_"):
                    continue
                config[k] = v
            config["tokenizer_vocab_size"] = tokenizer.vocab_size

            model = SymmetricTransformer2DModel(**config)
            model.load_state_dict(orig_model.state_dict(), strict=False)

            del orig_model
            gc.collect()
            torch.cuda.empty_cache()
    else:
        raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")

    model = torch.compile(model)

    if args.use_lora:
        lora_config = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            target_modules=args.lora_target_modules,
        )
        model.add_adapter(lora_config)

    model.train()
    for n, p in model.named_parameters():
        if "text_decoder" in n:
            p.requires_grad_(True)
        else:
            p.requires_grad_(False)

    if args.gradient_checkpointing:
        model.enable_gradient_checkpointing()
        if args.train_text_encoder:
            if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
                text_encoder[0].gradient_checkpointing_enable()
                text_encoder[1].gradient_checkpointing_enable()
            else:
                text_encoder.gradient_checkpointing_enable()

    if args.use_ema: # Not verify the robostness of this part
        ema = EMAModel(
            model.parameters(),
            decay=args.ema_decay,
            update_after_step=args.ema_update_after_step,
            model_cls= Transformer2DModel, 
            model_config=model.config,
        )

    def save_model_hook(models, weights, output_dir):
        if accelerator.is_main_process:
            transformer_lora_layers_to_save = None
            text_encoder_lora_layers_to_save = None

            for model_ in models:
                if isinstance(model_, type(accelerator.unwrap_model(model))):
                    if args.use_lora:
                        transformer_lora_layers_to_save = get_peft_model_state_dict(model_)
                    else:
                        model_.save_pretrained(os.path.join(output_dir, "transformer"))
                elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
                    if args.text_encoder_use_lora:
                        text_encoder_lora_layers_to_save = get_peft_model_state_dict(model_)
                    else:
                        model_.save_pretrained(os.path.join(output_dir, "text_encoder"))
                else:
                    raise ValueError(f"unexpected save model: {model_.__class__}")

                # make sure to pop weight so that corresponding model is not saved again
                weights.pop()

            if transformer_lora_layers_to_save is not None or text_encoder_lora_layers_to_save is not None:
                LoraLoaderMixin.save_lora_weights(
                    output_dir,
                    unet_lora_layers=transformer_lora_layers_to_save,
                    text_encoder_lora_layers=text_encoder_lora_layers_to_save,
                )

            if args.use_ema:
                ema.save_pretrained(os.path.join(output_dir, "ema_model"))

    def load_model_hook(models, input_dir):
        transformer = None
        text_encoder_ = None

        # this part is added for keep consistency when add model.compile() in the model
        def adap_compile(ori_dict):#add '_orig_mod.' to each key
            new_dict = {}
            for k,v in ori_dict.items():
                new_dict['_orig_mod.'+k] = v
            return new_dict
            
        while len(models) > 0:
            model_ = models.pop()

            if isinstance(model_, type(accelerator.unwrap_model(model))):
                if args.use_lora:
                    transformer = model_
                else:
                    if args.pretrained_model_architecture == "Meissonic":
                        load_model = SymmetricTransformer2DModel.from_pretrained(os.path.join(input_dir, "transformer"), low_cpu_mem_usage=False, device_map=None)
                    else:
                        raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")
                    model_.load_state_dict(adap_compile(load_model.state_dict()))
                    del load_model
            elif isinstance(model_, type(accelerator.unwrap_model(text_encoder))):
                if args.text_encoder_use_lora:
                    text_encoder_ = model_
                else:
                    try:
                        load_model = CLIPTextModelWithProjection.from_pretrained(os.path.join(input_dir, "text_encoder"))
                        model_.load_state_dict(load_model.state_dict())
                        # print('finished loading text encoder!')
                    except:
                        print('Not found text-encoder model in current folder. So we download one text encoder from Internet.')
                        load_model = CLIPTextModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K")
                        model_.load_state_dict(load_model.state_dict())
                    del load_model
            else:
                raise ValueError(f"unexpected save model: {model.__class__}")

        if transformer is not None or text_encoder_ is not None:
            lora_state_dict, network_alphas = LoraLoaderMixin.lora_state_dict(input_dir)
            LoraLoaderMixin.load_lora_into_text_encoder(
                lora_state_dict, network_alphas=network_alphas, text_encoder=text_encoder_
            )
            LoraLoaderMixin.load_lora_into_transformer(
                lora_state_dict, network_alphas=network_alphas, transformer=transformer
            )

        if args.use_ema:
            load_from = EMAModel.from_pretrained(os.path.join(input_dir, "ema_model"), model_cls=Transformer2DModel)
            ema.load_state_dict(adap_compile(load_from.state_dict()))
            del load_from

    accelerator.register_load_state_pre_hook(load_model_hook)
    accelerator.register_save_state_pre_hook(save_model_hook)

    if args.scale_lr:
        args.learning_rate = (
            args.learning_rate * args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
        )

    if args.use_8bit_adam:
        try:
            import bitsandbytes as bnb
        except ImportError:
            raise ImportError(
                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
            )

        optimizer_cls = bnb.optim.AdamW8bit
    else:
        optimizer_cls = torch.optim.AdamW


    optimizer_parameters = [p for p in model.parameters() if p.requires_grad]

    optimizer = optimizer_cls(
        optimizer_parameters,
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    logger.info("Creating dataloaders and lr_scheduler")

    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

    if args.instance_dataset == "MyParquetDataset":
        dataset = ImageCaptionDataset(
            root_dir=args.instance_data_dir,  # something like '../parquets_father_dir/'
            tokenizer=tokenizer,
            size=args.resolution,
            text_encoder_architecture=args.text_encoder_architecture
        )
    elif args.instance_dataset == 'HuggingFaceDataset':  # you can try this first, just download dataset from huggingface
        dataset = HuggingFaceDataset(
            hf_dataset=load_dataset(args.instance_data_dir, split="train"),  # something like './parquets_father_dir/'
            tokenizer=tokenizer,
            image_key='image',
            prompt_key='caption',
            prompt_prefix=args.prompt_prefix,
            size=args.resolution,
            text_encoder_architecture=args.text_encoder_architecture
        )
    elif args.instance_dataset == "DATA_TYPE":
        raise NotImplementedError("DATA_TYPE is not yet supported")
    else:
        assert False

    train_dataloader = DataLoader(
        dataset,
        batch_size=args.train_batch_size,
        shuffle=True,
        num_workers=args.dataloader_num_workers,
        collate_fn=default_collate,
        pin_memory=True,
    )
    train_dataloader.num_batches = len(train_dataloader)

    lr_scheduler = diffusers.optimization.get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_training_steps=args.max_train_steps * accelerator.num_processes,
        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
    )

    logger.info("Preparing model, optimizer and dataloaders")

    if args.train_text_encoder:
        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
            model, optimizer, lr_scheduler, train_dataloader, text_encoder[0], text_encoder[1] = accelerator.prepare(
                model, optimizer, lr_scheduler, train_dataloader, text_encoder[0], text_encoder[1]
            )
        else:
            model, optimizer, lr_scheduler, train_dataloader, text_encoder = accelerator.prepare(
            model, optimizer, lr_scheduler, train_dataloader, text_encoder
        )
    else:
        model, optimizer, lr_scheduler, train_dataloader = accelerator.prepare(
            model, optimizer, lr_scheduler, train_dataloader
        )

    train_dataloader.num_batches = len(train_dataloader)

    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    if not args.train_text_encoder:
        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
            text_encoder[0].to(device=accelerator.device, dtype=weight_dtype)
            text_encoder[1].to(device=accelerator.device, dtype=weight_dtype)
        else:
            text_encoder.to(device=accelerator.device, dtype=weight_dtype)

    vq_model.to(device=accelerator.device)

    if args.use_ema:
        ema.to(accelerator.device)

    with nullcontext() if args.train_text_encoder else torch.no_grad():
        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
            _input_ids_tmp_ = tokenize_prompt(tokenizer, "", args.text_encoder_architecture)
            _input_ids_tmp_[0] = _input_ids_tmp_[0].to(accelerator.device, non_blocking=True)
            _input_ids_tmp_[1] = _input_ids_tmp_[1].to(accelerator.device, non_blocking=True)
            empty_embeds, empty_clip_embeds = encode_prompt(
                text_encoder, _input_ids_tmp_, args.text_encoder_architecture
            )
        else:
            empty_embeds, empty_clip_embeds = encode_prompt(
                text_encoder, tokenize_prompt(tokenizer, "", args.text_encoder_architecture).to(accelerator.device, non_blocking=True), args.text_encoder_architecture
            )

        # There is a single image, we can just pre-encode the single prompt
        if args.instance_data_image is not None:
            prompt = os.path.splitext(os.path.basename(args.instance_data_image))[0]
            if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
                _input_ids_tmp_ = tokenize_prompt(tokenizer, prompt, args.text_encoder_architecture)
                _input_ids_tmp_[0] = _input_ids_tmp_[0].to(accelerator.device, non_blocking=True)
                _input_ids_tmp_[1] = _input_ids_tmp_[1].to(accelerator.device, non_blocking=True)
                empty_embeds, empty_clip_embeds = encode_prompt(
                    text_encoder, _input_ids_tmp_, args.text_encoder_architecture
                )
            else:
                encoder_hidden_states, cond_embeds = encode_prompt(
                    text_encoder, tokenize_prompt(tokenizer, prompt, args.text_encoder_architecture).to(accelerator.device, non_blocking=True), args.text_encoder_architecture
                )
            encoder_hidden_states = encoder_hidden_states.repeat(args.train_batch_size, 1, 1)
            cond_embeds = cond_embeds.repeat(args.train_batch_size, 1)

    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(train_dataloader.num_batches / args.gradient_accumulation_steps)
    # Afterwards we recalculate our number of training epochs.
    # Note: We are not doing epoch based training here, but just using this for book keeping and being able to
    # reuse the same training loop with other datasets/loaders.
    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    # Train!
    logger.info("***** Running training *****")
    logger.info(f"  Num training steps = {args.max_train_steps}")
    logger.info(f"  Instantaneous batch size per device = { args.train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")

    resume_from_checkpoint = args.resume_from_checkpoint
    if resume_from_checkpoint:
        if resume_from_checkpoint == "latest":
            # Get the most recent checkpoint
            dirs = os.listdir(args.output_dir)
            dirs = [d for d in dirs if d.startswith("checkpoint")]
            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
            if len(dirs) > 0:
                resume_from_checkpoint = os.path.join(args.output_dir, dirs[-1])
            else:
                resume_from_checkpoint = None

        if resume_from_checkpoint is None:
            accelerator.print(
                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
            )
        else:
            accelerator.print(f"Resuming from checkpoint {resume_from_checkpoint}")

    if resume_from_checkpoint is None:
        global_step = 0
        first_epoch = 0
    else:
        accelerator.load_state(resume_from_checkpoint)
        global_step = int(os.path.basename(resume_from_checkpoint).split("-")[1])
        first_epoch = global_step // num_update_steps_per_epoch

    # This is to solve the inconsistent tensor device issue
    if args.use_ema:
        ema.shadow_params = [p.to(accelerator.device) for p in ema.shadow_params]

    # As stated above, we are not doing epoch based training here, but just using this for book keeping and being able to
    # reuse the same training loop with other datasets/loaders.
    for epoch in range(first_epoch, num_train_epochs):
        for batch in train_dataloader:
            torch.cuda.empty_cache()
            with torch.no_grad():
                micro_conds = batch["micro_conds"].to(accelerator.device, non_blocking=True)
                pixel_values = batch["image"].to(accelerator.device, non_blocking=True)

                batch_size = pixel_values.shape[0]

                split_batch_size = args.split_vae_encode if args.split_vae_encode is not None else batch_size
                num_splits = math.ceil(batch_size / split_batch_size)
                image_tokens = []
                for i in range(num_splits):
                    start_idx = i * split_batch_size
                    end_idx = min((i + 1) * split_batch_size, batch_size)
                    bs = pixel_values.shape[0]
                    image_tokens.append(
                        vq_model.quantize(
                            vq_model.encode(
                                pixel_values[start_idx: end_idx]
                            ).latents
                        )[2][2].reshape(split_batch_size, -1)
                    )
                image_tokens = torch.cat(image_tokens, dim=0)

                batch_size, seq_len = image_tokens.shape

                timesteps = torch.ones(batch_size, device=image_tokens.device)
                mask_prob = torch.cos(timesteps * math.pi * 0.5)
                mask_prob = mask_prob.clip(args.min_masking_rate)

                num_token_masked = (seq_len * mask_prob).round().clamp(min=1)
                batch_randperm = torch.rand(batch_size, seq_len, device=image_tokens.device).argsort(dim=-1)
                mask = batch_randperm < num_token_masked.unsqueeze(-1)

                mask_id = accelerator.unwrap_model(model).config.vocab_size - 1
                input_ids = torch.where(mask, mask_id, image_tokens)
                # labels = torch.where(mask, image_tokens, -100)

                if "prompt_input_ids" in batch:
                    prompt_input_ids = batch["prompt_input_ids"]
                    labels = prompt_input_ids
                    with nullcontext() if args.train_text_encoder else torch.no_grad():
                        if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
                            batch["prompt_input_ids"][0] = batch["prompt_input_ids"][0].to(accelerator.device, non_blocking=True)
                            batch["prompt_input_ids"][1] = batch["prompt_input_ids"][1].to(accelerator.device, non_blocking=True)
                            encoder_hidden_states, cond_embeds = encode_prompt(
                                text_encoder, batch["prompt_input_ids"], args.text_encoder_architecture
                            )
                        else:
                            encoder_hidden_states, cond_embeds = encode_prompt(
                                text_encoder, batch["prompt_input_ids"].to(accelerator.device, non_blocking=True), args.text_encoder_architecture
                            )

                if args.cond_dropout_prob > 0.0:
                    assert encoder_hidden_states is not None

                    batch_size = encoder_hidden_states.shape[0]

                    mask = (
                        torch.zeros((batch_size, 1, 1), device=encoder_hidden_states.device).float().uniform_(0, 1)
                        < args.cond_dropout_prob
                    )

                    empty_embeds_ = empty_embeds.expand(batch_size, -1, -1)
                    encoder_hidden_states = torch.where(
                        (encoder_hidden_states * mask).bool(), encoder_hidden_states, empty_embeds_
                    )

                    empty_clip_embeds_ = empty_clip_embeds.expand(batch_size, -1)
                    cond_embeds = torch.where((cond_embeds * mask.squeeze(-1)).bool(), cond_embeds, empty_clip_embeds_)

                bs = input_ids.shape[0]
                vae_scale_factor = 2 ** (len(vq_model.config.block_out_channels) - 1)
                resolution = args.resolution // vae_scale_factor
                input_ids = input_ids.reshape(bs, resolution, resolution)

            # Train Step
            with accelerator.accumulate(model):
                if args.pretrained_model_architecture == 'Meissonic':
                   
                    if args.resolution == 1024: # only stage 3 and stage 4 do not apply 2*
                        img_ids = _prepare_latent_image_ids(input_ids.shape[0], input_ids.shape[-2], input_ids.shape[-1], input_ids.device, input_ids.dtype)
                    else:
                        img_ids = _prepare_latent_image_ids(input_ids.shape[0], 2 * input_ids.shape[-2], 2 * input_ids.shape[-1], input_ids.device, input_ids.dtype)

                    txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(device = input_ids.device, dtype = input_ids.dtype)
                   
                    logits = model(
                        hidden_states=input_ids, # should be (batch size, channel, height, width)
                        encoder_hidden_states=encoder_hidden_states, # should be (batch size, sequence_len, embed_dims)
                        micro_conds=micro_conds, # 
                        pooled_projections=cond_embeds, # should be (batch_size, projection_dim)
                        img_ids=img_ids,
                        txt_ids=txt_ids,
                        timestep=mask_prob * 1000,
                    )[1]
                    # print(logits.shape)
                    logits = logits.reshape(-1, tokenizer.vocab_size)
                    
                else:
                    raise ValueError(f"Unknown model architecture: {args.pretrained_model_architecture}")

                loss = F.cross_entropy(
                    logits,
                    labels.view(-1),
                    ignore_index=-100,
                    reduction="mean",
                )

                # Gather the losses across all processes for logging (if we use distributed training).
                avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
                avg_masking_rate = accelerator.gather(mask_prob.repeat(args.train_batch_size)).mean()

                accelerator.backward(loss)

                if args.max_grad_norm is not None and accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                optimizer.step()
                lr_scheduler.step()

                optimizer.zero_grad(set_to_none=True)

            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                if args.use_ema:
                    ema.step(model.parameters())

                if (global_step + 1) % args.logging_steps == 0:
                    logs = {
                        "step_loss": avg_loss.item(),
                        "lr": lr_scheduler.get_last_lr()[0],
                        "avg_masking_rate": avg_masking_rate.item(),
                    }
                    accelerator.log(logs, step=global_step + 1)

                    logger.info(
                        f"Step: {global_step + 1} "
                        f"Loss: {avg_loss.item():0.4f} "
                        f"LR: {lr_scheduler.get_last_lr()[0]:0.6f}"
                    )

                if (global_step + 1) % args.checkpointing_steps == 0:
                    save_checkpoint(args, accelerator, global_step + 1, logger)

                if (global_step + 1) % args.validation_steps == 0 and accelerator.is_main_process:
                    if args.use_ema:
                        ema.store(model.parameters())
                        ema.copy_to(model.parameters())

                    with torch.no_grad():
                        logger.info("Generating images...")

                        model.eval()

                        scheduler = Scheduler.from_pretrained(
                            args.pretrained_model_name_or_path,
                            subfolder="scheduler",
                            revision=args.revision,
                            variant=args.variant,
                        )

                        pipe = UnifiedPipeline(
                            transformer=accelerator.unwrap_model(model),
                            tokenizer=tokenizer,
                            text_encoder=text_encoder,
                            vqvae=vq_model,
                            scheduler=scheduler,
                        )

                        output = pipe(
                            prompt=args.validation_prompts, 
                            height=args.resolution, 
                            width=args.resolution, 
                            guidance_scale=9, 
                            num_inference_steps=64
                        )
                        pil_images = output.images
                        prompts = output.prompts
                        print(prompts)

                        result=[]
                        for img in pil_images:
                            if not isinstance(img, torch.Tensor):
                                img = transforms.ToTensor()(img)
                            result.append(img.unsqueeze(0))
                        result = torch.cat(result,dim=0)
                        result = make_grid(result, nrow=3)
                        save_image(result,os.path.join(args.output_dir, str(global_step)+'_text2image_1024_CFG-9.png'))
                        
                        # 保存为JSON
                        output_data = {
                            "step": global_step,
                            "prompts": prompts,
                            "images": [f"{global_step}_text2image_1024_CFG-9_{i}.png" for i in range(len(pil_images))]
                        }

                        with open(os.path.join(args.output_dir, f"prompts_{global_step}.json"), "w") as f:
                            json.dump(output_data, f, indent=2)
                        
                        model.train()

                        if args.train_text_encoder:
                            if args.text_encoder_architecture == "CLIP_T5_base": # Not support yet. Only support open_clip
                                text_encoder[0].train()
                                text_encoder[1].trian()
                            else:
                                text_encoder.train()

                    if args.use_ema:
                        ema.restore(model.parameters())

                global_step += 1

            # Stop training if max steps is reached
            if global_step >= args.max_train_steps:
                break
        # End for

    accelerator.wait_for_everyone()

    # Evaluate and save checkpoint at the end of training
    save_checkpoint(args, accelerator, global_step, logger)

    # Save the final trained checkpoint
    if accelerator.is_main_process:
        model = accelerator.unwrap_model(model)
        if args.use_ema:
            ema.copy_to(model.parameters())
        model.save_pretrained(args.output_dir)

    accelerator.end_training()


if __name__ == "__main__":
    main(parse_args())