|
|
import sys |
|
|
|
|
|
import torch |
|
|
import os |
|
|
import json |
|
|
import argparse |
|
|
|
|
|
sys.path.append(os.getcwd()) |
|
|
|
|
|
from diffusers import StableDiffusion3Pipeline, FluxPipeline, FlowMatchHeunDiscreteScheduler, FlowMatchEulerDiscreteScheduler |
|
|
from scheduler.scheduling_flow_match_euler_discrete_lm import FlowMatchEulerDiscreteLMScheduler |
|
|
from tqdm import tqdm |
|
|
|
|
|
def main(): |
|
|
parser = argparse.ArgumentParser(description="sampling script for T2I-Bench.") |
|
|
parser.add_argument('--test_num', type=int, default=10) |
|
|
parser.add_argument('--start_index', type=int, default=0) |
|
|
parser.add_argument('--num_inference_steps', type=int, default=10) |
|
|
parser.add_argument('--guidance', type=float, default=7.5) |
|
|
parser.add_argument('--sampler_type', type = str, default='fm_euler') |
|
|
parser.add_argument('--model_id', type=str, default='XXX') |
|
|
parser.add_argument('--save_dir', type=str, default='results/') |
|
|
parser.add_argument('--lamb', type=float, default=5.0) |
|
|
parser.add_argument('--kappa', type=float, default=0.0) |
|
|
parser.add_argument('--freeze', type=float, default=0.0) |
|
|
parser.add_argument('--dataset_category', type=str, default="color") |
|
|
parser.add_argument('--dataset_path', type=str, default="T2I-CompBench-main") |
|
|
parser.add_argument('--dtype', type=str, default='bf16') |
|
|
parser.add_argument('--device', type=str, default='cuda') |
|
|
|
|
|
args = parser.parse_args() |
|
|
dtype = None |
|
|
if args.dtype in ['fp32']: |
|
|
dtype = torch.float32 |
|
|
elif args.dtype in ['fp64']: |
|
|
dtype = torch.float64 |
|
|
elif args.dtype in ['fp16']: |
|
|
dtype = torch.float16 |
|
|
elif args.dtype in ['bf16']: |
|
|
dtype = torch.bfloat16 |
|
|
|
|
|
start_index = args.start_index |
|
|
sampler_type = args.sampler_type |
|
|
test_num = args.test_num |
|
|
guidance_scale = args.guidance |
|
|
num_inference_steps = args.num_inference_steps |
|
|
lamb = args.lamb |
|
|
freeze = args.freeze |
|
|
kappa = args.kappa |
|
|
model_id = args.model_id |
|
|
device = args.device |
|
|
|
|
|
|
|
|
sd_pipe = FluxPipeline.from_pretrained( |
|
|
model_id, |
|
|
torch_dtype=dtype, safety_checker=None) |
|
|
sd_pipe = sd_pipe.to(device) |
|
|
print("flux model loaded") |
|
|
|
|
|
if sampler_type in ['fm_euler']: |
|
|
pass |
|
|
elif sampler_type in ['lml_euler']: |
|
|
sd_pipe.scheduler = FlowMatchEulerDiscreteLMScheduler.from_config(sd_pipe.scheduler.config) |
|
|
sd_pipe.scheduler.lamb = lamb |
|
|
sd_pipe.scheduler.lm = True |
|
|
sd_pipe.scheduler.kappa = kappa |
|
|
else: |
|
|
raise ValueError(f"invalid: '{sampler_type}'.") |
|
|
|
|
|
save_dir = args.save_dir |
|
|
|
|
|
if sampler_type in ['lml_euler']: |
|
|
save_dir = os.path.join(save_dir, "flux", args.dataset_category, sampler_type + "_lamda_" + str(lamb)) |
|
|
else: |
|
|
save_dir = os.path.join(save_dir, "flux", args.dataset_category, sampler_type) |
|
|
|
|
|
save_dir = os.path.join(save_dir, "samples") |
|
|
if not os.path.exists(save_dir): |
|
|
os.makedirs(save_dir, exist_ok=True) |
|
|
|
|
|
def getT2IDataset(file_path): |
|
|
with open(file_path, "r", encoding="utf-8") as file: |
|
|
for line in file: |
|
|
stripped_line = line.strip() |
|
|
if stripped_line: |
|
|
yield stripped_line |
|
|
|
|
|
|
|
|
dataset_path = os.path.join(args.dataset_path, 'examples/dataset', args.dataset_category + '_val.txt') |
|
|
count = 0 |
|
|
with tqdm(total=300 * test_num, desc="Generating Images") as pbar: |
|
|
try: |
|
|
for prompt in getT2IDataset(dataset_path): |
|
|
for seed in range(start_index, start_index + test_num): |
|
|
torch.manual_seed(seed) |
|
|
res = sd_pipe(prompt=prompt, num_inference_steps=num_inference_steps, |
|
|
guidance_scale=guidance_scale, generator=None, width=512, height=512).images[0] |
|
|
res.save(os.path.join(save_dir, f"{prompt}_{count:06d}.png")) |
|
|
count += 1 |
|
|
pbar.update(1) |
|
|
except FileNotFoundError: |
|
|
print(f"dataset can not be found: {dataset_path}") |
|
|
except Exception as e: |
|
|
print(f"unknown error: {str(e)}") |
|
|
print(f"{dataset_path} finish") |
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
