import os import argparse import random import socket import json import time import math from tqdm import tqdm import numpy as np import torch import torch.distributed as dist from torchvision.transforms.functional import to_pil_image from diffusers.models import AutoencoderKL from transformers import AutoTokenizer, AutoModel import models from transport import ODE # Adapted from pipelines.StableDiffusionXLPipeline.encode_prompt def encode_prompt(prompt_batch, text_encoder, tokenizer, proportion_empty_prompts, is_train=True): captions = [] for caption in prompt_batch: if random.random() < proportion_empty_prompts: captions.append("") elif isinstance(caption, str): captions.append(caption) elif isinstance(caption, (list, np.ndarray)): # take a random caption if there are multiple captions.append(random.choice(caption) if is_train else caption[0]) with torch.no_grad(): text_inputs = tokenizer( captions, padding=True, pad_to_multiple_of=8, max_length=256, truncation=True, return_tensors="pt", ) text_input_ids = text_inputs.input_ids prompt_masks = text_inputs.attention_mask prompt_embeds = text_encoder(input_ids=text_input_ids.cuda(),attention_mask=prompt_masks.cuda(),output_hidden_states=True,).hidden_states[-2] return prompt_embeds, prompt_masks def none_or_str(value): if value == 'None': return None return value def main(args, rank, master_port): # Setup PyTorch: torch.set_grad_enabled(False) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(args.num_gpus) os.environ["MASTER_PORT"] = str(master_port) os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["TOKENIZERS_PARALLELISM"] = "false" dist.init_process_group("nccl") torch.cuda.set_device(rank) train_args = torch.load(os.path.join(args.ckpt, "model_args.pth")) if dist.get_rank() == 0: print("Loaded model arguments:", json.dumps(train_args.__dict__, indent=2)) if dist.get_rank() == 0: print(f"Creating lm: Gemma-2B") dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[args.precision] tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b") tokenizer.padding_side = 'right' text_encoder = AutoModel.from_pretrained( "google/gemma-2b", torch_dtype=dtype, ).eval().cuda() # Load scheduler and models cap_feat_dim = text_encoder.config.hidden_size if dist.get_rank() == 0: print(f"Creating vae: {train_args.vae}") vae = AutoencoderKL.from_pretrained( (f"stabilityai/sd-vae-ft-{train_args.vae}" if train_args.vae != "sdxl" else "stabilityai/sdxl-vae"), torch_dtype=torch.float32, ).cuda() if dist.get_rank() == 0: print(f"Creating DiT: {train_args.model}") # latent_size = train_args.image_size // 8 model = models.__dict__[train_args.model]( qk_norm=train_args.qk_norm, cap_feat_dim=cap_feat_dim, ) model.eval().to("cuda", dtype=dtype) if args.debug == False: # assert train_args.model_parallel_size == args.num_gpus if args.ema: print("Loading ema model.") ckpt = torch.load(os.path.join( args.ckpt, f"consolidated{'_ema' if args.ema else ''}." f"{rank:02d}-of-{args.num_gpus:02d}.pth", ), map_location="cpu") model.load_state_dict(ckpt, strict=True) sample_folder_dir = args.image_save_path if rank == 0: os.makedirs(sample_folder_dir, exist_ok=True) os.makedirs(os.path.join(sample_folder_dir, 'images'), exist_ok=True) print(f"Saving .png samples at {sample_folder_dir}") dist.barrier() info_path = os.path.join(args.image_save_path, 'data.json') if os.path.exists(info_path): with open(info_path, 'r') as f: info = json.loads(f.read()) collected_id = [] for i in info: collected_id.append(f'{id(i["caption"])}_{i["resolution"]}') else: info = [] collected_id = [] captions = [] with open(args.caption_path, 'r', encoding='utf-8') as file: for line in file: text = line.strip() if text: captions.append(line.strip()) total = len(info) resolution = args.resolution with torch.autocast("cuda", dtype): for res in resolution: for idx, caption in tqdm(enumerate(captions)): if int(args.seed) != 0: torch.random.manual_seed(int(args.seed)) sample_id = f'{idx}_{res.split(":")[-1]}' if sample_id in collected_id: continue caps_list = [caption] res_cat, resolution = res.split(":") res_cat = int(res_cat) do_extrapolation = res_cat > 1024 n = len(caps_list) w, h = resolution.split("x") w, h = int(w), int(h) latent_w, latent_h = w // 8, h // 8 z = torch.randn([1, 4, latent_w, latent_h], device="cuda").to(dtype) z = z.repeat(n * 2, 1, 1, 1) with torch.no_grad(): cap_feats, cap_mask = encode_prompt([caps_list] + [""], text_encoder, tokenizer, 0.0) cap_mask = cap_mask.to(cap_feats.device) model_kwargs = dict( cap_feats=cap_feats, cap_mask=cap_mask, cfg_scale=args.cfg_scale, ) if args.proportional_attn: model_kwargs["proportional_attn"] = True model_kwargs["base_seqlen"] = (train_args.image_size // 16) ** 2 else: model_kwargs["proportional_attn"] = False model_kwargs["base_seqlen"] = None if do_extrapolation and args.scaling_method == "Time-aware": model_kwargs["scale_factor"] = math.sqrt(w * h / train_args.image_size**2) model_kwargs["scale_watershed"] = args.scaling_watershed else: model_kwargs["scale_factor"] = 1.0 model_kwargs["scale_watershed"] = 1.0 samples = ODE(args.num_sampling_steps, args.solver, args.time_shifting_factor).sample(z, model.forward_with_cfg, **model_kwargs)[-1] samples = samples[:1] factor = 0.18215 if train_args.vae != "sdxl" else 0.13025 samples = vae.decode(samples / factor).sample samples = (samples + 1.0) / 2.0 samples.clamp_(0.0, 1.0) # Save samples to disk as individual .png files for i, (sample, cap) in enumerate(zip(samples, caps_list)): img = to_pil_image(sample.float()) save_path = f"{args.image_save_path}/images/{args.solver}_{args.num_sampling_steps}_{sample_id}.png" img.save(save_path) info.append({ 'caption': cap, 'image_url': f"{args.image_save_path}/images/{args.solver}_{args.num_sampling_steps}_{sample_id}.png", 'resolution': f'res: {resolution}\ntime_shift: {args.time_shifting_factor}', 'solver': args.solver, 'num_sampling_steps': args.num_sampling_steps }) with open(info_path, 'w') as f: f.write(json.dumps(info)) total += len(samples) dist.barrier() dist.barrier() dist.barrier() dist.destroy_process_group() def find_free_port() -> int: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(("", 0)) port = sock.getsockname()[1] sock.close() return port if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--cfg_scale", type=float, default=4.0) parser.add_argument("--num_sampling_steps", type=int, default=250) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--ckpt", type=str, required=True) parser.add_argument("--solver", type=str, default="euler") parser.add_argument( "--precision", type=str, choices=["fp32", "bf16"], default="bf16", ) parser.add_argument("--num_gpus", type=int, default=1) parser.add_argument("--ema", action="store_true", help="Use EMA models.") parser.set_defaults(ema=True) parser.add_argument( "--image_save_path", type=str, default='samples', help="If specified, overrides the default image save path " "(sample{_ema}.png in the model checkpoint directory)." ) parser.add_argument( "--time_shifting_factor", type=float, default=1.0, ) parser.add_argument( "--caption_path", type=str, default='prompts.txt', ) parser.add_argument( "--resolution", type=str, default='', nargs="+", ) parser.add_argument( "--tokenizer_path", type=str, default='', ) parser.add_argument( "--proportional_attn", type=bool, default=True ) parser.add_argument( "--scaling_method", type=str, default="Time-aware", ) parser.add_argument( "--scaling_watershed", type=float, default=0.3, ) parser.add_argument("--debug", action="store_true") parser.add_argument("--batch_size", type=int, default=8) args = parser.parse_known_args()[0] master_port = find_free_port() assert args.num_gpus == 1, "Multi-GPU sampling is currently not supported." main(args, 0, master_port)