# Authors: Hui Ren (rhfeiyang.github.io) import torch from PIL import Image import argparse import os, json, random import matplotlib.pyplot as plt import glob, re from tqdm import tqdm import numpy as np import sys import gc from transformers import CLIPTextModel, CLIPTokenizer, BertModel, BertTokenizer # import train_util from utils.train_util import get_noisy_image, encode_prompts from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel, LMSDiscreteScheduler, DDIMScheduler, PNDMScheduler from typing import Any, Dict, List, Optional, Tuple, Union from utils.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV import argparse # from diffusers.training_utils import EMAModel import shutil import yaml from easydict import EasyDict from utils.metrics import StyleContentMetric from torchvision import transforms from custom_datasets.coco import CustomCocoCaptions from custom_datasets.imagepair import ImageSet from custom_datasets import get_dataset # from stable_diffusion.utils.modules import get_diffusion_modules # from diffusers import StableDiffusionImg2ImgPipeline from diffusers.utils.torch_utils import randn_tensor import pickle import time def flush(): torch.cuda.empty_cache() gc.collect() def get_train_method(lora_weight): if lora_weight is None: return 'None' if 'full' in lora_weight: train_method = 'full' elif "down_1_up_2_attn" in lora_weight: train_method = 'up_2_attn' print(f"Using up_2_attn for {lora_weight}") elif "down_2_up_1_up_2_attn" in lora_weight: train_method = 'down_2_up_2_attn' elif "down_2_up_2_attn" in lora_weight: train_method = 'down_2_up_2_attn' elif "down_2_attn" in lora_weight: train_method = 'down_2_attn' elif 'noxattn' in lora_weight: train_method = 'noxattn' elif "xattn" in lora_weight: train_method = 'xattn' elif "attn" in lora_weight: train_method = 'attn' elif "all_up" in lora_weight: train_method = 'all_up' else: train_method = 'None' return train_method def get_validation_dataloader(infer_prompts:list[str]=None, infer_images :list[str]=None,resolution=512, batch_size=10, num_workers=4, val_set="laion_pop500"): data_transforms = transforms.Compose( [ transforms.Resize(resolution), transforms.CenterCrop(resolution), ] ) def preprocess(example): ret={} ret["image"] = data_transforms(example["image"]) if "image" in example else None if "caption" in example: if isinstance(example["caption"][0], list): ret["caption"] = example["caption"][0][0] else: ret["caption"] = example["caption"][0] if "seed" in example: ret["seed"] = example["seed"] if "id" in example: ret["id"] = example["id"] if "path" in example: ret["path"] = example["path"] return ret def collate_fn(examples): out = {} if "image" in examples[0]: pixel_values = [example["image"] for example in examples] out["pixel_values"] = pixel_values # notice: only take the first prompt for each image if "caption" in examples[0]: prompts = [example["caption"] for example in examples] out["prompts"] = prompts if "seed" in examples[0]: seeds = [example["seed"] for example in examples] out["seed"] = seeds if "path" in examples[0]: paths = [example["path"] for example in examples] out["path"] = paths return out if infer_prompts is None: if val_set == "lhq500": dataset = get_dataset("lhq_sub500", get_val=False)["train"] elif val_set == "custom_coco100": dataset = get_dataset("custom_coco100", get_val=False)["train"] elif val_set == "custom_coco500": dataset = get_dataset("custom_coco500", get_val=False)["train"] elif os.path.isdir(val_set): image_folder = os.path.join(val_set, "paintings") caption_folder = os.path.join(val_set, "captions") dataset = ImageSet(folder=image_folder, caption=caption_folder, keep_in_mem=True) elif "custom_caption" in val_set: from custom_datasets.custom_caption import Caption_set name = val_set.replace("custom_caption_", "") dataset = Caption_set(set_name = name) elif val_set == "laion_pop500": dataset = get_dataset("laion_pop500", get_val=False)["train"] elif val_set == "laion_pop500_first_sentence": dataset = get_dataset("laion_pop500_first_sentence", get_val=False)["train"] else: raise ValueError("Unknown dataset") dataset.with_transform(preprocess) elif isinstance(infer_prompts, torch.utils.data.Dataset): dataset = infer_prompts try: dataset.with_transform(preprocess) except: pass else: class Dataset(torch.utils.data.Dataset): def __init__(self, prompts, images=None): self.prompts = prompts self.images = images self.get_img = False if images is not None: assert len(prompts) == len(images) self.get_img = True if isinstance(images[0], str): self.images = [Image.open(image).convert("RGB") for image in images] else: self.images = [None] * len(prompts) def __len__(self): return len(self.prompts) def __getitem__(self, idx): img = self.images[idx] if self.get_img and img is not None: img = data_transforms(img) return {"caption": self.prompts[idx], "image":img} dataset = Dataset(infer_prompts, infer_images) dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn, drop_last=False, num_workers=num_workers, pin_memory=True) return dataloader def get_lora_network(unet , lora_path, train_method="None", rank=1, alpha=1.0, device="cuda", weight_dtype=torch.float32): if train_method in [None, "None"]: train_method = get_train_method(lora_path) print(f"Train method: {train_method}") network_type = "c3lier" if train_method == 'xattn': network_type = 'lierla' modules = DEFAULT_TARGET_REPLACE if network_type == "c3lier": modules += UNET_TARGET_REPLACE_MODULE_CONV alpha = 1 if "rank" in lora_path: rank = int(re.search(r'rank(\d+)', lora_path).group(1)) if 'alpha1' in lora_path: alpha = 1.0 print(f"Rank: {rank}, Alpha: {alpha}") network = LoRANetwork( unet, rank=rank, multiplier=1.0, alpha=alpha, train_method=train_method, ).to(device, dtype=weight_dtype) if lora_path not in [None, "None"]: lora_state_dict = torch.load(lora_path) miss = network.load_state_dict(lora_state_dict, strict=False) print(f"Missing: {miss}") ret = {"network": network, "train_method": train_method} return ret def get_model(pretrained_ckpt_path, unet_ckpt=None,revision=None, variant=None, lora_path=None, weight_dtype=torch.float32, device="cuda"): modules = {} pipe = DiffusionPipeline.from_pretrained(pretrained_ckpt_path, revision=revision, variant=variant) if unet_ckpt is not None: pipe.unet.from_pretrained(unet_ckpt, subfolder="unet_ema", revision=revision, variant=variant) unet = pipe.unet vae = pipe.vae text_encoder = pipe.text_encoder tokenizer = pipe.tokenizer modules["unet"] = unet modules["vae"] = vae modules["text_encoder"] = text_encoder modules["tokenizer"] = tokenizer # tokenizer = modules["tokenizer"] unet.enable_xformers_memory_efficient_attention() unet.to(device, dtype=weight_dtype) if weight_dtype != torch.bfloat16: vae.to(device, dtype=torch.float32) else: vae.to(device, dtype=weight_dtype) text_encoder.to(device, dtype=weight_dtype) if lora_path is not None: network = get_lora_network(unet, lora_path, device=device, weight_dtype=weight_dtype) modules["network"] = network return modules @torch.no_grad() def inference(network: LoRANetwork, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, vae: AutoencoderKL, unet: UNet2DConditionModel, noise_scheduler: LMSDiscreteScheduler, dataloader, height:int, width:int, scales:list = np.linspace(0,2,5),save_dir:str=None, seed:int = None, weight_dtype: torch.dtype = torch.float32, device: torch.device="cuda", batch_size:int=1, steps:int=50, guidance_scale:float=7.5, start_noise:int=800, uncond_prompt:str=None, uncond_embed=None, style_prompt = None, show:bool = False, no_load:bool=False, from_scratch=False): print(f"save dir: {save_dir}") if start_noise < 0: assert from_scratch network = network.eval() unet = unet.eval() vae = vae.eval() do_convert = not from_scratch if not do_convert: try: dataloader.dataset.get_img = False except: pass scales = list(scales) else: scales = ["Real Image"] + list(scales) if not no_load and os.path.exists(os.path.join(save_dir, "infer_imgs.pickle")): with open(os.path.join(save_dir, "infer_imgs.pickle"), 'rb') as f: pred_images = pickle.load(f) take=True for key in scales: if key not in pred_images: take=False break if take: print(f"Found existing inference results in {save_dir}", flush=True) return pred_images max_length = tokenizer.model_max_length pred_images = {scale :[] for scale in scales} all_seeds = {scale:[] for scale in scales} prompts = [] ori_prompts = [] if save_dir is not None: img_output_dir = os.path.join(save_dir, "outputs") os.makedirs(img_output_dir, exist_ok=True) if uncond_embed is None: if uncond_prompt is None: uncond_input_text = [""] else: uncond_input_text = [uncond_prompt] uncond_embed = encode_prompts(tokenizer = tokenizer, text_encoder = text_encoder, prompts = uncond_input_text) for batch in dataloader: ori_prompt = batch["prompts"] image = batch["pixel_values"] if do_convert else None if do_convert: pred_images["Real Image"] += image if isinstance(ori_prompt, list): if isinstance(text_encoder, CLIPTextModel): # trunc prompts for clip encoder ori_prompt = [p.split(".")[0]+"." for p in ori_prompt] prompt = [f"{p.strip()[::-1].replace('.', '',1)[::-1]} in the style of {style_prompt}" for p in ori_prompt] if style_prompt is not None else ori_prompt else: if isinstance(text_encoder, CLIPTextModel): ori_prompt = ori_prompt.split(".")[0]+"." prompt = f"{prompt.strip()[::-1].replace('.', '',1)[::-1]} in the style of {style_prompt}" if style_prompt is not None else ori_prompt bcz = len(prompt) single_seed = seed if dataloader.batch_size == 1 and seed is None: if "seed" in batch: single_seed = batch["seed"][0] print(f"{prompt}, seed={single_seed}") # text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt").to(device) # original_embeddings = text_encoder(**text_input)[0] prompts += prompt ori_prompts += ori_prompt # style_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt").to(device) # # style_embeddings = text_encoder(**style_input)[0] # style_embeddings = text_encoder(style_input.input_ids, return_dict=False)[0] style_embeddings = encode_prompts(tokenizer = tokenizer, text_encoder = text_encoder, prompts = prompt) original_embeddings = encode_prompts(tokenizer = tokenizer, text_encoder = text_encoder, prompts = ori_prompt) if uncond_embed.shape[0] == 1 and bcz > 1: uncond_embeddings = uncond_embed.repeat(bcz, 1, 1) else: uncond_embeddings = uncond_embed style_text_embeddings = torch.cat([uncond_embeddings, style_embeddings]).to(weight_dtype) # original_embeddings = torch.cat([uncond_embeddings, original_embeddings]).to(weight_dtype) generator = torch.manual_seed(single_seed) if single_seed is not None else None noise_scheduler.set_timesteps(steps) if do_convert: noised_latent, _, _ = get_noisy_image(image, vae, generator, unet, noise_scheduler, total_timesteps=int((1000-start_noise)/1000 *steps)) else: latent_shape = (bcz, 4, height//8, width//8) noised_latent = randn_tensor(latent_shape, generator=generator, device=vae.device) noised_latent = noised_latent.to(unet.dtype) noised_latent = noised_latent * noise_scheduler.init_noise_sigma for scale in scales: start_time = time.time() if not isinstance(scale, float) and not isinstance(scale, int): continue latents = noised_latent.clone().to(weight_dtype).to(device) noise_scheduler.set_timesteps(steps) for t in tqdm(noise_scheduler.timesteps): if do_convert and t>start_noise: continue else: if t > start_noise and start_noise >= 0: current_scale = 0 else: current_scale = scale network.set_lora_slider(scale=current_scale) text_embedding = style_text_embeddings # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. latent_model_input = torch.cat([latents] * 2).to(weight_dtype) latent_model_input = noise_scheduler.scale_model_input(latent_model_input, timestep=t).to(weight_dtype) # predict the noise residual with network: # print(f"dtype: {latent_model_input.dtype}, {text_embedding.dtype}, t={t}") noise_pred = unet(latent_model_input, t , encoder_hidden_states=text_embedding).sample # perform guidance noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) # compute the previous noisy sample x_t -> x_t-1 if isinstance(noise_scheduler, DDPMScheduler): latents = noise_scheduler.step(noise_pred, t, latents, generator=torch.manual_seed(single_seed+t) if single_seed is not None else None).prev_sample else: latents = noise_scheduler.step(noise_pred, t, latents).prev_sample # scale and decode the image latents with vae latents = 1 / 0.18215 * latents.to(vae.dtype) with torch.no_grad(): image = vae.decode(latents).sample image = (image / 2 + 0.5).clamp(0, 1) image = image.detach().cpu().permute(0, 2, 3, 1).to(torch.float32).numpy() images = (image * 255).round().astype("uint8") pil_images = [Image.fromarray(image) for image in images] pred_images[scale]+=pil_images all_seeds[scale] += [single_seed] * bcz end_time = time.time() print(f"Time taken for one batch, Art Adapter scale={scale}: {end_time-start_time}", flush=True) if save_dir is not None or show: end_idx = len(list(pred_images.values())[0]) for i in range(end_idx-bcz, end_idx): keys = list(pred_images.keys()) images_list = [pred_images[key][i] for key in keys] prompt = prompts[i] if len(scales)==1: plt.imshow(images_list[0]) plt.axis('off') plt.title(f"{prompt}_{single_seed}_start{start_noise}", fontsize=20) else: fig, ax = plt.subplots(1, len(images_list), figsize=(len(scales)*5,6), layout="constrained") for id, a in enumerate(ax): a.imshow(images_list[id]) if isinstance(scales[id], float) or isinstance(scales[id], int): a.set_title(f"Art Adapter scale={scales[id]}", fontsize=20) else: a.set_title(f"{keys[id]}", fontsize=20) a.axis('off') # plt.suptitle(f"{os.path.basename(lora_weight).replace('.pt','')}", fontsize=20) # plt.tight_layout() # if do_convert: # plt.suptitle(f"{prompt}\nseed{single_seed}_start{start_noise}_guidance{guidance_scale}", fontsize=20) # else: # plt.suptitle(f"{prompt}\nseed{single_seed}_from_scratch_guidance{guidance_scale}", fontsize=20) if save_dir is not None: plt.savefig(f"{img_output_dir}/{prompt.replace(' ', '_')[:100]}_seed{single_seed}_start{start_noise}.png") if show: plt.show() plt.close() flush() if save_dir is not None: with open(os.path.join(save_dir, "infer_imgs.pickle" ), 'wb') as f: pickle.dump(pred_images, f) with open(os.path.join(save_dir, "all_seeds.pickle"), 'wb') as f: to_save={"all_seeds":all_seeds, "batch_size":batch_size} pickle.dump(to_save, f) for scale, images in pred_images.items(): subfolder = os.path.join(save_dir,"images", f"{scale}") os.makedirs(subfolder, exist_ok=True) used_prompt = ori_prompts if (isinstance(scale, float) or isinstance(scale, int)): #and scale != 0: used_prompt = prompts for i, image in enumerate(images): if scale == "Real Image": suffix = "" else: suffix = f"_seed{all_seeds[scale][i]}" image.save(os.path.join(subfolder, f"{used_prompt[i].replace(' ', '_')[:100]}{suffix}.jpg")) with open(os.path.join(save_dir, "infer_prompts.txt"), 'w') as f: for prompt in prompts: f.write(f"{prompt}\n") with open(os.path.join(save_dir, "ori_prompts.txt"), 'w') as f: for prompt in ori_prompts: f.write(f"{prompt}\n") print(f"Saved inference results to {save_dir}", flush=True) return pred_images, prompts @torch.no_grad() def infer_metric(ref_image_folder,pred_images, prompts, save_dir, start_noise=""): prompts = [prompt.split(" in the style of ")[0] for prompt in prompts] scores = {} original_images = pred_images["Real Image"] if "Real Image" in pred_images else None metric = StyleContentMetric(ref_image_folder) for scale, images in pred_images.items(): score = metric(images, original_images, prompts) scores[scale] = score print(f"Style transfer score at scale {scale}: {score}") scores["ref_path"] = ref_image_folder save_name = f"scores_start{start_noise}.json" os.makedirs(save_dir, exist_ok=True) with open(os.path.join(save_dir, save_name), 'w') as f: json.dump(scores, f, indent=2) return scores def parse_args(): parser = argparse.ArgumentParser(description='Inference with LoRA') parser.add_argument('--lora_weights', type=str, default=["None"], nargs='+', help='path to your model file') parser.add_argument('--prompts', type=str, default=[], nargs='+', help='prompts to try') parser.add_argument("--prompt_file", type=str, default=None, help="path to the prompt file") parser.add_argument("--prompt_file_key", type=str, default="prompts", help="key to the prompt file") parser.add_argument('--resolution', type=int, default=512, help='resolution of the image') parser.add_argument('--seed', type=int, default=None, help='seed for the random number generator') parser.add_argument("--start_noise", type=int, default=800, help="start noise") parser.add_argument("--from_scratch", default=False, action="store_true", help="from scratch") parser.add_argument("--ref_image_folder", type=str, default=None, help="folder containing reference images") parser.add_argument("--show", action="store_true", help="show the image") parser.add_argument("--batch_size", type=int, default=1, help="batch size") parser.add_argument("--scales", type=float, default=[0.,1.], nargs='+', help="scales to test") parser.add_argument("--train_method", type=str, default=None, help="train method") # parser.add_argument("--vae_path", type=str, default="CompVis/stable-diffusion-v1-4", help="Path to the VAE model.") # parser.add_argument("--text_encoder_path", type=str, default="CompVis/stable-diffusion-v1-4", help="Path to the text encoder model.") parser.add_argument("--pretrained_model_name_or_path", type=str, default="rhfeiyang/art-free-diffusion-v1", help="Path to the pretrained model.") parser.add_argument("--unet_ckpt", default=None, type=str, help="Path to the unet checkpoint") parser.add_argument("--guidance_scale", type=float, default=5.0, help="guidance scale") parser.add_argument("--infer_mode", default="sks_art", help="inference mode") #, choices=["style", "ori", "artist", "sks_art","Peter"] parser.add_argument("--save_dir", type=str, default="inference_output", help="save directory") parser.add_argument("--num_workers", type=int, default=4, help="number of workers") parser.add_argument("--no_load", action="store_true", help="no load the pre-inferred results") parser.add_argument("--infer_prompts", type=str, default=None, nargs="+", help="prompts to infer") parser.add_argument("--infer_images", type=str, default=None, nargs="+", help="images to infer") parser.add_argument("--rank", type=int, default=1, help="rank of the lora") parser.add_argument("--val_set", type=str, default="laion_pop500", help="validation set") parser.add_argument("--folder_name", type=str, default=None, help="folder name") parser.add_argument("--scheduler_type",type=str, choices=["ddpm", "ddim", "pndm","lms"], default="ddpm", help="scheduler type") parser.add_argument("--infer_steps", type=int, default=50, help="inference steps") parser.add_argument("--weight_dtype", type=str, default="fp32", help="weight dtype") parser.add_argument("--custom_coco_cap", action="store_true", help="use custom coco caption") args = parser.parse_args() if args.infer_prompts is not None and len(args.infer_prompts) == 1 and os.path.isfile(args.infer_prompts[0]): if args.infer_prompts[0].endswith(".txt") and args.custom_coco_cap: args.infer_prompts = CustomCocoCaptions(custom_file=args.infer_prompts[0]) elif args.infer_prompts[0].endswith(".txt"): with open(args.infer_prompts[0], 'r') as f: args.infer_prompts = f.readlines() args.infer_prompts = [prompt.strip() for prompt in args.infer_prompts] elif args.infer_prompts[0].endswith(".csv"): from custom_datasets.custom_caption import Caption_set caption_set = Caption_set(args.infer_prompts[0]) args.infer_prompts = caption_set if args.infer_mode == "style": with open(os.path.join(args.ref_image_folder, "style_label.txt"), 'r') as f: args.style_label = f.readlines()[0].strip() elif args.infer_mode == "artist": with open(os.path.join(args.ref_image_folder, "style_label.txt"), 'r') as f: args.style_label = f.readlines()[0].strip() args.style_label = args.style_label.split(",")[0].strip() elif args.infer_mode == "ori": args.style_label = None else: args.style_label = args.infer_mode.replace("_", " ") if args.ref_image_folder is not None: args.ref_image_folder = os.path.join(args.ref_image_folder, "paintings") if args.start_noise < 0: args.from_scratch = True print(args.__dict__) return args def main(args): lora_weights = args.lora_weights if len(lora_weights) == 1 and isinstance(lora_weights[0], str) and os.path.isdir(lora_weights[0]): lora_weights = glob.glob(os.path.join(lora_weights[0], "*.pt")) lora_weights=sorted(lora_weights, reverse=True) width = args.resolution height = args.resolution steps = args.infer_steps revision = None device = 'cuda' rank = args.rank if args.weight_dtype == "fp32": weight_dtype = torch.float32 elif args.weight_dtype=="fp16": weight_dtype = torch.float16 elif args.weight_dtype=="bf16": weight_dtype = torch.bfloat16 modules = get_model(args.pretrained_model_name_or_path, unet_ckpt=args.unet_ckpt, revision=revision, variant=None, lora_path=None, weight_dtype=weight_dtype, device=device, ) if args.scheduler_type == "pndm": noise_scheduler = PNDMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler") elif args.scheduler_type == "ddpm": noise_scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler") elif args.scheduler_type == "ddim": noise_scheduler = DDIMScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False, prediction_type="epsilon", ) elif args.scheduler_type == "lms": noise_scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) else: raise ValueError("Unknown scheduler type") cache=EasyDict() cache.modules = modules unet = modules["unet"] vae = modules["vae"] text_encoder = modules["text_encoder"] tokenizer = modules["tokenizer"] unet.requires_grad_(False) # Move unet, vae and text_encoder to device and cast to weight_dtype vae.requires_grad_(False) text_encoder.requires_grad_(False) ## dataloader dataloader = get_validation_dataloader(infer_prompts=args.infer_prompts, infer_images=args.infer_images, resolution=args.resolution, batch_size=args.batch_size, num_workers=args.num_workers, val_set=args.val_set) for lora_weight in lora_weights: print(f"Testing {lora_weight}") # for different seeds on same prompt seed = args.seed network_ret = get_lora_network(unet, lora_weight, train_method=args.train_method, rank=rank, alpha=1.0, device=device, weight_dtype=weight_dtype) network = network_ret["network"] train_method = network_ret["train_method"] if args.save_dir is not None: save_dir = args.save_dir if args.style_label is not None: save_dir = os.path.join(save_dir, f"{args.style_label.replace(' ', '_')}") else: save_dir = os.path.join(save_dir, f"ori/{args.start_noise}") else: if args.folder_name is not None: folder_name = args.folder_name else: folder_name = "validation" if args.infer_prompts is None else "validation_prompts" save_dir = os.path.join(os.path.dirname(lora_weight), f"{folder_name}/{train_method}", os.path.basename(lora_weight).replace('.pt','').split('_')[-1]) if args.infer_prompts is None: save_dir = os.path.join(save_dir, f"{args.val_set}") infer_config = f"{args.scheduler_type}{args.infer_steps}_{args.weight_dtype}_guidance{args.guidance_scale}" save_dir = os.path.join(save_dir, infer_config) os.makedirs(save_dir, exist_ok=True) if args.from_scratch: save_dir = os.path.join(save_dir, "from_scratch") else: save_dir = os.path.join(save_dir, "transfer") save_dir = os.path.join(save_dir, f"start{args.start_noise}") os.makedirs(save_dir, exist_ok=True) with open(os.path.join(save_dir, "infer_args.yaml"), 'w') as f: yaml.dump(vars(args), f) # save code code_dir = os.path.join(save_dir, "code") os.makedirs(code_dir, exist_ok=True) current_file = os.path.basename(__file__) shutil.copy(__file__, os.path.join(code_dir, current_file)) with torch.no_grad(): pred_images, prompts = inference(network, tokenizer, text_encoder, vae, unet, noise_scheduler, dataloader, height, width, args.scales, save_dir, seed, weight_dtype, device, args.batch_size, steps, guidance_scale=args.guidance_scale, start_noise=args.start_noise, show=args.show, style_prompt=args.style_label, no_load=args.no_load, from_scratch=args.from_scratch) if args.ref_image_folder is not None: flush() print("Calculating metrics") infer_metric(args.ref_image_folder, pred_images, save_dir, args.start_noise) if __name__ == "__main__": args = parse_args() main(args)