260 lines
9.0 KiB
Python
260 lines
9.0 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
import socket
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
from tqdm import tqdm
|
|
from PIL import Image
|
|
|
|
from data.prefix_instruction import test_task_dicts
|
|
from data.data_reader import T2IItemProcessor
|
|
from data.data_utils import check_item_graph200k
|
|
from visualcloze import VisualClozeModel
|
|
|
|
|
|
def concat_images_grid(images):
|
|
row_images = []
|
|
for row in images:
|
|
row_widths, row_heights = zip(*(img.size for img in row))
|
|
total_width = sum(row_widths)
|
|
max_height = max(row_heights)
|
|
|
|
new_row_img = Image.new('RGB', (total_width, max_height), (255, 255, 255))
|
|
|
|
x_offset = 0
|
|
for img in row:
|
|
new_row_img.paste(img, (x_offset, 0))
|
|
x_offset += img.width
|
|
|
|
row_images.append(new_row_img)
|
|
|
|
col_widths, col_heights = zip(*(img.size for img in row_images))
|
|
total_height = sum(col_heights)
|
|
max_width = max(col_widths)
|
|
|
|
final_img = Image.new('RGB', (max_width, total_height), (255, 255, 255))
|
|
|
|
y_offset = 0
|
|
for img in row_images:
|
|
final_img.paste(img, (0, y_offset))
|
|
y_offset += img.height
|
|
|
|
return final_img
|
|
|
|
|
|
def main(args, rank, master_port):
|
|
# Setup PyTorch:
|
|
torch.set_grad_enabled(False)
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
torch.cuda.manual_seed_all(args.seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = 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")
|
|
|
|
print("Init model")
|
|
visualcloze = VisualClozeModel(
|
|
model_path=args.model_path,
|
|
resolution=args.resolution,
|
|
lora_rank=args.lora_rank
|
|
)
|
|
|
|
item_processor = T2IItemProcessor(None, resolution=args.resolution)
|
|
|
|
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'{i["idx"]}')
|
|
else:
|
|
info = []
|
|
collected_id = []
|
|
|
|
data = []
|
|
# Loading data for testing
|
|
with open(args.data_path, "r", encoding="utf-8") as file:
|
|
if args.data_path.endswith('.jsonl'):
|
|
for line in file:
|
|
data.append(json.loads(line))
|
|
else:
|
|
data = json.load(file)
|
|
|
|
for idx, item in tqdm(enumerate(data[::1])):
|
|
|
|
for context_num in [1, 2, 3]:
|
|
for task_type in test_task_dicts:
|
|
for image_type_list in task_type["image_list"]:
|
|
if not check_item_graph200k(item, image_type_list):
|
|
continue
|
|
task_name = "_".join(image_type_list)
|
|
|
|
context_ids = []
|
|
context_items = []
|
|
|
|
while len(context_items) < context_num - 1:
|
|
next_idx = random.randint(0, len(data) - 1)
|
|
if next_idx == idx or next_idx in context_ids:
|
|
continue
|
|
if not check_item_graph200k(data[next_idx], image_type_list):
|
|
continue
|
|
context_ids.append(next_idx)
|
|
context_items.append(data[next_idx])
|
|
|
|
all_items = context_items + [item]
|
|
_, images, prompts, _, grid_shape = item_processor.graph200k_process_item(all_items, image_type_list, context_num, training_mode=False)
|
|
grid_w, grid_h = grid_shape
|
|
if f'{idx}' in collected_id:
|
|
continue
|
|
|
|
images[-1][-1] = None
|
|
visualcloze.set_grid_size(grid_h, grid_w)
|
|
ret = visualcloze.process_images(
|
|
images, prompts,
|
|
seed=args.seed,
|
|
cfg=args.guidance_scale,
|
|
steps=args.num_sampling_steps,
|
|
upsampling_noise=None,
|
|
upsampling_steps=None,
|
|
is_upsampling=False
|
|
)[-1]
|
|
images[-1][-1] = ret
|
|
|
|
image = concat_images_grid(images)
|
|
save_path = f"{args.image_save_path}/images/{args.solver}_{args.num_sampling_steps}_{idx}_{context_num}_{task_name}.jpg"
|
|
image.save(save_path, format='JPEG', quality=95)
|
|
|
|
info.append(
|
|
{
|
|
"idx": idx,
|
|
"caption": prompts,
|
|
"image_url": f"{args.image_save_path}/images/{args.solver}_{args.num_sampling_steps}_{idx}_{context_num}_{task_name}.jpg",
|
|
"seed": args.seed,
|
|
"guidance_scale": args.guidance_scale,
|
|
"steps": args.num_sampling_steps,
|
|
"solver": args.solver,
|
|
"num_sampling_steps": args.num_sampling_steps,
|
|
}
|
|
)
|
|
|
|
with open(info_path, "w") as f:
|
|
f.write(json.dumps(info))
|
|
|
|
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("--text_encoder", type=str, nargs='+', default=['t5', 'clip'], help="List of text encoders to use (e.g., t5, clip, gemma)")
|
|
parser.add_argument("--num_sampling_steps", type=int, default=250)
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--model_path", 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(
|
|
"--data_path",
|
|
type=str,
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--resolution",
|
|
type=int,
|
|
default=384,
|
|
)
|
|
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(
|
|
"--atol",
|
|
type=float,
|
|
default=1e-6,
|
|
help="Absolute tolerance for the ODE solver.",
|
|
)
|
|
parser.add_argument(
|
|
"--rtol",
|
|
type=float,
|
|
default=1e-3,
|
|
help="Relative tolerance for the ODE solver.",
|
|
)
|
|
parser.add_argument("--do_shift", default=True)
|
|
parser.add_argument("--attn_token_select", action="store_true")
|
|
parser.add_argument("--mlp_token_select", action="store_true")
|
|
parser.add_argument("--reverse", action="store_true", help="run the ODE solver in reverse.")
|
|
parser.add_argument("--use_flash_attn", type=bool, default=True)
|
|
parser.add_argument("--debug", action="store_true")
|
|
parser.add_argument("--batch_size", type=int, default=8)
|
|
parser.add_argument("--max_length", type=int, default=512, help="Max length for T5.")
|
|
parser.add_argument("--root_path", type=str, default="")
|
|
parser.add_argument("--guidance_scale", type=float, default=30.0)
|
|
parser.add_argument("--do_classifier_free_guidance", action="store_true")
|
|
parser.add_argument("--lora_rank", type=int, default=128)
|
|
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) |