import torch torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.set_float32_matmul_precision('high') setattr(torch.nn.Linear, 'reset_parameters', lambda self: None) # disable default parameter init for faster speed setattr(torch.nn.LayerNorm, 'reset_parameters', lambda self: None) # disable default parameter init for faster speed import torch.nn.functional as F import torch.distributed as dist import os import math import json import argparse import pandas as pd from tqdm import tqdm from PIL import Image from tokenizer.tokenizer_image.vq_model import VQ_models from language.t5 import T5Embedder from autoregressive.models.gpt import GPT_models from autoregressive.models.generate import generate os.environ["TOKENIZERS_PARALLELISM"] = "false" def main(args): # Setup PyTorch: assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage" torch.set_grad_enabled(False) # Setup DDP: dist.init_process_group("nccl") rank = dist.get_rank() device = rank % torch.cuda.device_count() seed = args.global_seed * dist.get_world_size() + rank torch.manual_seed(seed) torch.cuda.set_device(device) print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.") # create and load model vq_model = VQ_models[args.vq_model]( codebook_size=args.codebook_size, codebook_embed_dim=args.codebook_embed_dim) vq_model.to(device) vq_model.eval() checkpoint = torch.load(args.vq_ckpt, map_location="cpu") vq_model.load_state_dict(checkpoint["model"]) del checkpoint print(f"image tokenizer is loaded") # create and load gpt model precision = {'none': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}[args.precision] latent_size = args.image_size // args.downsample_size gpt_model = GPT_models[args.gpt_model]( block_size=latent_size ** 2, cls_token_num=args.cls_token_num, model_type=args.gpt_type, ).to(device=device, dtype=precision) checkpoint = torch.load(args.gpt_ckpt, map_location="cpu") if "model" in checkpoint: # ddp model_weight = checkpoint["model"] elif "module" in checkpoint: # deepspeed model_weight = checkpoint["module"] elif "state_dict" in checkpoint: model_weight = checkpoint["state_dict"] else: raise Exception("please check model weight") gpt_model.load_state_dict(model_weight, strict=False) gpt_model.eval() del checkpoint print(f"gpt model is loaded") if args.compile: print(f"compiling the model...") gpt_model = torch.compile( gpt_model, mode="reduce-overhead", fullgraph=True ) # requires PyTorch 2.0 (optional) else: print(f"no need to compile model in demo") assert os.path.exists(args.t5_path) t5_model = T5Embedder( device=device, local_cache=True, cache_dir=args.t5_path, dir_or_name=args.t5_model_type, torch_dtype=precision, model_max_length=args.t5_feature_max_len, ) print(f"t5 model is loaded") # Create folder to save samples: model_string_name = args.gpt_model.replace("/", "-") ckpt_string_name = os.path.basename(args.gpt_ckpt).replace(".pth", "").replace(".pt", "") prompt_name = args.prompt_csv.split('/')[-1].split('.')[0].lower() folder_name = f"{model_string_name}-{ckpt_string_name}-{prompt_name}-size-{args.image_size}-size-{args.image_size}-{args.vq_model}-" \ f"topk-{args.top_k}-topp-{args.top_p}-temperature-{args.temperature}-" \ f"cfg-{args.cfg_scale}-seed-{args.global_seed}" sample_folder_dir = f"{args.sample_dir}/{folder_name}" if rank == 0: os.makedirs(f"{sample_folder_dir}/images", exist_ok=True) print(f"Saving .png samples at {sample_folder_dir}/images") dist.barrier() df = pd.read_csv(args.prompt_csv, delimiter='\t') prompt_list = df['Prompt'].tolist() # Figure out how many samples we need to generate on each GPU and how many iterations we need to run: n = args.per_proc_batch_size global_batch_size = n * dist.get_world_size() num_fid_samples = min(args.num_fid_samples, len(prompt_list)) # To make things evenly-divisible, we'll sample a bit more than we need and then discard the extra samples: total_samples = int(math.ceil(num_fid_samples / global_batch_size) * global_batch_size) if rank == 0: print(f"Total number of images that will be sampled: {total_samples}") assert total_samples % dist.get_world_size() == 0, "total_samples must be divisible by world_size" samples_needed_this_gpu = int(total_samples // dist.get_world_size()) assert samples_needed_this_gpu % n == 0, "samples_needed_this_gpu must be divisible by the per-GPU batch size" iterations = int(samples_needed_this_gpu // n) pbar = range(iterations) pbar = tqdm(pbar) if rank == 0 else pbar total = 0 for _ in pbar: # Select text prompt prompt_batch = [] for i in range(n): index = i * dist.get_world_size() + rank + total prompt_batch.append(prompt_list[index] if index < len(prompt_list) else "a cute dog") # Sample inputs: caption_embs, emb_masks = t5_model.get_text_embeddings(prompt_batch) if not args.no_left_padding: new_emb_masks = torch.flip(emb_masks, dims=[-1]) new_caption_embs = [] for idx, (caption_emb, emb_mask) in enumerate(zip(caption_embs, emb_masks)): valid_num = int(emb_mask.sum().item()) # prompt_cur = prompt_batch[idx] # print(f' prompt {idx} token len: {valid_num} : {prompt_cur}') new_caption_emb = torch.cat([caption_emb[valid_num:], caption_emb[:valid_num]]) new_caption_embs.append(new_caption_emb) new_caption_embs = torch.stack(new_caption_embs) else: new_caption_embs, new_emb_masks = caption_embs, emb_masks c_indices = new_caption_embs * new_emb_masks[:,:, None] c_emb_masks = new_emb_masks qzshape = [len(c_indices), args.codebook_embed_dim, latent_size, latent_size] index_sample = generate( gpt_model, c_indices, latent_size ** 2, c_emb_masks, cfg_scale=args.cfg_scale, temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, sample_logits=True, ) samples = vq_model.decode_code(index_sample, qzshape) # output value is between [-1, 1] samples = torch.clamp(127.5 * samples + 128.0, 0, 255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy() # Save samples to disk as individual .png files for i, sample in enumerate(samples): index = i * dist.get_world_size() + rank + total Image.fromarray(sample).save(f"{sample_folder_dir}/images/{index:06d}.png") total += global_batch_size # Make sure all processes have finished saving their samples before attempting to convert to .npz dist.barrier() if rank == 0: # Save infer result in a jsonl file json_items = [] for idx, prompt in enumerate(prompt_list): image_path = os.path.join(sample_folder_dir, "images", f"{idx:06d}.png") json_items.append({"text": prompt, "image_path": image_path}) res_jsonl_path = os.path.join(sample_folder_dir, "result.jsonl") print(f"Save jsonl to {res_jsonl_path}...") with open(res_jsonl_path, "w") as f: for item in json_items: f.write(json.dumps(item) + "\n") # Save captions to txt caption_path = os.path.join(sample_folder_dir, "captions.txt") print(f"Save captions to {caption_path}...") with open(caption_path, "w") as f: for item in prompt_list: f.write(f"{item}\n") print("Done.") dist.barrier() dist.destroy_process_group() if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--prompt-csv", type=str, default='evaluations/t2i/PartiPrompts.tsv') parser.add_argument("--t5-path", type=str, default='pretrained_models/t5-ckpt') parser.add_argument("--t5-model-type", type=str, default='flan-t5-xl') parser.add_argument("--t5-feature-max-len", type=int, default=120) parser.add_argument("--t5-feature-dim", type=int, default=2048) parser.add_argument("--no-left-padding", action='store_true', default=False) parser.add_argument("--gpt-model", type=str, choices=list(GPT_models.keys()), default="GPT-XL") parser.add_argument("--gpt-ckpt", type=str, default=None) parser.add_argument("--gpt-type", type=str, choices=['c2i', 't2i'], default="t2i", help="class->image or text->image") parser.add_argument("--cls-token-num", type=int, default=120, help="max token number of condition input") parser.add_argument("--precision", type=str, default='bf16', choices=["none", "fp16", "bf16"]) parser.add_argument("--compile", action='store_true', default=False) parser.add_argument("--vq-model", type=str, choices=list(VQ_models.keys()), default="VQ-16") parser.add_argument("--vq-ckpt", type=str, default=None, help="ckpt path for vq model") parser.add_argument("--codebook-size", type=int, default=16384, help="codebook size for vector quantization") parser.add_argument("--codebook-embed-dim", type=int, default=8, help="codebook dimension for vector quantization") parser.add_argument("--image-size", type=int, choices=[256, 384, 512], default=512) parser.add_argument("--downsample-size", type=int, choices=[8, 16], default=16) parser.add_argument("--num-classes", type=int, default=1000) parser.add_argument("--cfg-scale", type=float, default=7.5) parser.add_argument("--sample-dir", type=str, default="samples_parti", help="samples_coco or samples_parti") parser.add_argument("--per-proc-batch-size", type=int, default=32) parser.add_argument("--num-fid-samples", type=int, default=30000) parser.add_argument("--global-seed", type=int, default=0) parser.add_argument("--top-k", type=int, default=1000, help="top-k value to sample with") parser.add_argument("--temperature", type=float, default=1.0, help="temperature value to sample with") parser.add_argument("--top-p", type=float, default=1.0, help="top-p value to sample with") args = parser.parse_args() main(args)