From 0e9ab058925827f1c5346e4433f16fdb16b095cd Mon Sep 17 00:00:00 2001 From: Fereshteh Forghani Date: Thu, 22 Feb 2024 16:59:18 -0500 Subject: [PATCH 1/3] Add data generation script --- img2img.py | 172 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) create mode 100644 img2img.py diff --git a/img2img.py b/img2img.py new file mode 100644 index 0000000..5284300 --- /dev/null +++ b/img2img.py @@ -0,0 +1,172 @@ +import argparse +import torch +from diffusers import StableUnCLIPImg2ImgPipeline, DDIMScheduler, DPMSolverSinglestepScheduler +# from xformers.ops import MemoryEfficientAttentionFlashAttentionOp +from torchvision import datasets, transforms +from torch.utils.data import Dataset, DataLoader +import time +import os +import json + + + +class StableGenerator(object): + def __init__(self, opt): + self.opt = opt + # model + self.model = StableUnCLIPImg2ImgPipeline.from_pretrained( + "stabilityai/stable-diffusion-2-1-unclip-small", torch_dtype=torch.float16, variation="fp16" + ) + + device = torch.device("cuda") if opt.device == "cuda" else torch.device("cpu") + self.model.to(device) + print(f"Using device: {device}") + + if opt.dpm: + self.model.scheduler = DPMSolverSinglestepScheduler.from_config(self.model.scheduler.config, rescale_betas_zero_snr=True) + else: + self.model.scheduler = DDIMScheduler.from_config(self.model.scheduler.config, rescale_betas_zero_snr=True, timestep_spacing="trailing") + + print("Scheduler:", self.model.scheduler) + + # self.model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp) + # self.model.vae.enable_xformers_memory_efficient_attention(attention_op=None) + + # image size + self.height = self.opt.img_save_size + self.width = self.opt.img_save_size + + # inference steps + self.num_inference_steps = self.opt.steps + + # eta (0, 1) + self.eta = self.opt.ddim_eta + + def generate(self, input_image, n_sample_per_image=10): + transfoem_2 = transforms.Resize(size=(self.height, self.width)) + synth_images = self.model(input_image, eta=self.eta, num_images_per_prompt=n_sample_per_image, num_inference_steps=self.num_inference_steps).images + synth_images = [transfoem_2(img) for img in synth_images] + return synth_images + + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--outdir", + type=str, + nargs="?", + help="dir to write results to", + default="outputs/txt2img-samples" + ) + parser.add_argument( + "--img_save_size", + type=int, + default=224, + help="image saving size" + ) + parser.add_argument( + "--steps", + type=int, + default=10, + help="number of ddim sampling steps", + ) + parser.add_argument( + "--dpm", + action='store_true', + help="use DPM (2) sampler", + ) + parser.add_argument( + "--ddim", + action='store_true', + help="use ddim sampler", + ) + parser.add_argument( + "--ddim_eta", + type=float, + default=0.0, + help="ddim eta (eta=0.0 corresponds to deterministic sampling", + ) + parser.add_argument( + "--n_samples", + type=int, + default=1, + help="how many samples to produce for each given prompt. A.k.a. batch size", + ) + parser.add_argument( + "--batch_size", + type=int, + default=8, + help="batch size", + ) + parser.add_argument( + "--device", + type=str, + default="cuda", + help="device to use for inference", + ) + parser.add_argument( + "--start", + type=int, + default=0, + help="start index", + ) + parser.add_argument( + "--end", + type=int, + help="end index", + ) + opt = parser.parse_args() + + # s = 1 + # size = 128 + # color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s) + + if opt.outdir is not None: + os.makedirs(opt.outdir, exist_ok=True) + + transform_list = [ + transforms.Resize(size=(opt.img_save_size,opt.img_save_size)), + # transforms.ToTensor(), + ] + transform = transforms.Compose(transform_list) + + imagenet_dataset = datasets.ImageNet("/local/ssd/m2kowal/imagenet", split="train", transform = transform) + # root_folder = "/local/ssd/m2kowal/imagenet/" + + Stable_generator = StableGenerator(opt) + n = len(imagenet_dataset) + assert opt.start < n + assert opt.end <= n + + for i in range(opt.start, opt.end): + print(f"Batch {i}") + batch = imagenet_dataset[i] + images = batch[0] + start = time.time() + generated_images = Stable_generator.generate(images, n_sample_per_image=opt.n_samples) + end = time.time() + print(f"Time taken to generate images: {end-start} seconds") + + ## save images + path = imagenet_dataset.samples[i][0] + start = time.time() + save_images(path, generated_images, opt.outdir) + end = time.time() + print(f"Time taken to save images: {end-start} seconds") + +def save_images(path, images, out_dir): + for j, img in enumerate(images): + out_folder = path.split("/")[-1].split(".")[0].split("_")[0] # get the class name + file_name = path.split("/")[-1].split(".")[0] # get the (class name_image number) + save_folder = os.path.join(out_dir, out_folder) # create a folder for each class + + if not os.path.exists(save_folder): + os.makedirs(save_folder, exist_ok=True) + + for j, img in enumerate(images): + save_file = os.path.join(save_folder, f"{file_name}_{j+1}.png") + img.save(save_file) + +if __name__ == "__main__": + main() \ No newline at end of file From 9c47af040b709ad446c418a29e0228f0fbf0df2c Mon Sep 17 00:00:00 2001 From: Vahid Reza Khazaie Date: Fri, 23 Feb 2024 15:10:49 -0500 Subject: [PATCH 2/3] add icgan generation script --- create_imagenet_icgan.py | 108 +++++++++++++++++++++++++++++++++++++++ icgan/config.py | 11 ++-- icgan/icgan_inference.py | 10 ++-- train_simclr.slrm | 13 ++--- 4 files changed, 127 insertions(+), 15 deletions(-) create mode 100644 create_imagenet_icgan.py diff --git a/create_imagenet_icgan.py b/create_imagenet_icgan.py new file mode 100644 index 0000000..c179660 --- /dev/null +++ b/create_imagenet_icgan.py @@ -0,0 +1,108 @@ +import argparse +from torchvision import datasets, transforms +import time +import os +from icgan.config import get_config +from icgan.icgan_inference import ICGANInference +from pytorch_pretrained_biggan import convert_to_images +from icgan.data_utils import utils as data_utils +import torch + +import numpy as np + +def save(out, torch_format=True): + if torch_format: + with torch.no_grad(): + out = out.cpu().numpy() + img = convert_to_images(out)[0] + return img + +def save_images(path, images, out_dir): + for j, img in enumerate(images): + out_folder = path.split("/")[-1].split(".")[0].split("_")[0] # get the class name + file_name = path.split("/")[-1].split(".")[0] # get the (class name_image number) + save_folder = os.path.join(out_dir, out_folder) # create a folder for each class + + if not os.path.exists(save_folder): + os.makedirs(save_folder, exist_ok=True) + + for j, img in enumerate(images): + save_file = os.path.join(save_folder, f"{file_name}_{j+1}.png") + img = save(img[np.newaxis, ...], torch_format=False) + img.save(save_file) + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--outdir", + type=str, + nargs="?", + help="dir to write results to", + default="/projects/imagenet_synthetic/synthetic_icgan" + ) + # parser.add_argument( + # "--img_save_size", + # type=int, + # default=224, + # help="image saving size" + # ) + # parser.add_argument( + # "--batch_size", + # type=int, + # default=1, + # help="batch size", + # ) + parser.add_argument( + "--start", + type=int, + default=0, + help="start index", + ) + parser.add_argument( + "--end", + type=int, + default=-1, + help="end index", + ) + args = parser.parse_args() + + if args.outdir is not None: + os.makedirs(args.outdir, exist_ok=True) + + # Initialize the ICGANInference class + config = get_config() + icgan_inference = ICGANInference(config) + + transform = transforms.Compose([ + data_utils.CenterCropLongEdge(), + # transforms.Resize((args.img_save_size, args.img_save_size)), + transforms.ToTensor(), + transforms.Normalize(config.norm_mean, config.norm_std)]) + + + imagenet_dataset = datasets.ImageNet("/scratch/ssd004/datasets/imagenet256", split="train", transform=transform) + + n = len(imagenet_dataset) + if args.end == -1: + args.end = n + assert args.start < n + assert args.end <= n + + for i in range(args.start, args.end): + print(f"Batch {i}") + batch = imagenet_dataset[i] + images = batch[0] + start = time.time() + generated_images = icgan_inference.run_inference(input_image_tensor=images.unsqueeze(0)) + end = time.time() + print(f"Time taken to generate images: {end-start} seconds") + + ## save images + path = imagenet_dataset.samples[i][0] + start = time.time() + save_images(path, generated_images, args.outdir) + end = time.time() + print(f"Time taken to save images: {end-start} seconds") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/icgan/config.py b/icgan/config.py index cd57edb..1599dc1 100644 --- a/icgan/config.py +++ b/icgan/config.py @@ -4,13 +4,16 @@ def get_config(): config = ml_collections.ConfigDict() - config.num_samples = 16 + config.seed = 123 + config.num_samples = 1 config.truncation = 0.7 config.stochastic_truncation = False config.noise_size = 128 - config.batch_size = 4 - config.experiment_name = "icgan_biggan_imagenet_res256" - config.feat_ext_path = "swav_pretrained.pth.tar" + config.batch_size = 1 + config.experiment_name = ( + "/ssd003/projects/aieng/genssl/icgan_biggan_imagenet_res256" + ) + config.feat_ext_path = "/ssd003/projects/aieng/genssl/swav_pretrained.pth.tar" config.size = 256 config.norm_mean = torch.Tensor([0.485, 0.456, 0.406]).view(3, 1, 1) config.norm_std = torch.Tensor([0.229, 0.224, 0.225]).view(3, 1, 1) diff --git a/icgan/icgan_inference.py b/icgan/icgan_inference.py index 700ee93..fbb046d 100644 --- a/icgan/icgan_inference.py +++ b/icgan/icgan_inference.py @@ -5,8 +5,8 @@ from scipy.stats import truncnorm from torch import nn import torchvision.transforms as transforms -import inference.utils as inference_utils -import data_utils.utils as data_utils +from icgan.inference import utils as inference_utils +from icgan.data_utils import utils as data_utils class ICGANInference: @@ -24,14 +24,14 @@ def replace_to_inplace_relu(self, model): else: self.replace_to_inplace_relu(child) - def load_icgan(self, root_="./"): + def load_icgan(self, root_=""): root = os.path.join(root_, self.config.experiment_name) config = torch.load("%s/%s.pth" % (root, "state_dict_best0"))["config"] config["weights_root"] = root_ config["model_backbone"] = "biggan" config["experiment_name"] = self.config.experiment_name G, config = inference_utils.load_model_inference(config) - G.cuda() + G = G.cuda() G.eval() return G @@ -62,7 +62,7 @@ def normality_loss(self, vec): return mu2 + sigma2 - torch.log(sigma2) - 1 def load_generative_model(self): - model = self.load_icgan(root_="./") + model = self.load_icgan(root_="") return model def load_feature_extractor(self, feat_ext_path): diff --git a/train_simclr.slrm b/train_simclr.slrm index 477a1fb..4021805 100644 --- a/train_simclr.slrm +++ b/train_simclr.slrm @@ -1,6 +1,6 @@ #!/bin/bash -#SBATCH --job-name=train_sunrgbd +#SBATCH --job-name=train_simclr #SBATCH --partition=t4v2 #SBATCH --time=12:00:00 #SBATCH --nodes=1 @@ -8,7 +8,7 @@ #SBATCH --ntasks-per-node=4 #SBATCH --cpus-per-task=4 #SBATCH --mem-per-cpu=2G -#SBATCH --output=slurm-%N-%j.out +#SBATCH --output=slurm2-%N-%j.out PY_ARGS=${@:1} @@ -16,10 +16,10 @@ PY_ARGS=${@:1} source /ssd003/projects/aieng/envs/genssl2/bin/activate export NCCL_ASYNC_ERROR_HANDLING=1 # set to 1 for NCCL backend -export CUDA_LAUNCH_BLOCKING=1 +# export CUDA_LAUNCH_BLOCKING=1 export MASTER_ADDR=$(hostname) -export MASTER_PORT=45679 +export MASTER_PORT=45678 export PYTHONPATH="." nvidia-smi @@ -28,5 +28,6 @@ nvidia-smi srun python run_simCLR.py \ --fp16-precision \ --distributed_mode \ ---batch-size=4 \ ---icgan_augmentation +--batch-size=16 \ +--icgan_augmentation \ +--subset_fraction 0.05 \ No newline at end of file From 8ff76a8202a996dd6acd88211bd444d1de01d9a9 Mon Sep 17 00:00:00 2001 From: Vahid Reza Khazaie Date: Fri, 23 Feb 2024 17:10:07 -0500 Subject: [PATCH 3/3] revise the script --- create_imagenet_icgan.py | 95 ++++++++++++++++++---------------------- 1 file changed, 43 insertions(+), 52 deletions(-) diff --git a/create_imagenet_icgan.py b/create_imagenet_icgan.py index c179660..9994e51 100644 --- a/create_imagenet_icgan.py +++ b/create_imagenet_icgan.py @@ -7,8 +7,43 @@ from pytorch_pretrained_biggan import convert_to_images from icgan.data_utils import utils as data_utils import torch - import numpy as np +from tqdm import tqdm + +parser = argparse.ArgumentParser() +parser.add_argument( + "--outdir", + type=str, + nargs="?", + help="dir to write results to", + default="/projects/imagenet_synthetic/synthetic_icgan" +) +parser.add_argument( + "--img_save_size", + type=int, + default=224, + help="image saving size" +) +parser.add_argument( + "--start", + type=int, + default=0, + help="start index", +) +parser.add_argument( + "--end", + type=int, + default=-1, + help="end index", +) +parser.add_argument( + "--ith_sample", + type=int, + default=0, + help="end index", +) + +args = parser.parse_args() def save(out, torch_format=True): if torch_format: @@ -18,7 +53,7 @@ def save(out, torch_format=True): return img def save_images(path, images, out_dir): - for j, img in enumerate(images): + for img in images: out_folder = path.split("/")[-1].split(".")[0].split("_")[0] # get the class name file_name = path.split("/")[-1].split(".")[0] # get the (class name_image number) save_folder = os.path.join(out_dir, out_folder) # create a folder for each class @@ -26,56 +61,20 @@ def save_images(path, images, out_dir): if not os.path.exists(save_folder): os.makedirs(save_folder, exist_ok=True) - for j, img in enumerate(images): - save_file = os.path.join(save_folder, f"{file_name}_{j+1}.png") - img = save(img[np.newaxis, ...], torch_format=False) - img.save(save_file) + save_file = os.path.join(save_folder, f"{file_name}_{args.ith_sample}.JPEG") + pil_img = save(img[np.newaxis, ...], torch_format=False) + pil_img.thumbnail((args.img_save_size, args.img_save_size)) + pil_img.save(save_file, format="JPEG") def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--outdir", - type=str, - nargs="?", - help="dir to write results to", - default="/projects/imagenet_synthetic/synthetic_icgan" - ) - # parser.add_argument( - # "--img_save_size", - # type=int, - # default=224, - # help="image saving size" - # ) - # parser.add_argument( - # "--batch_size", - # type=int, - # default=1, - # help="batch size", - # ) - parser.add_argument( - "--start", - type=int, - default=0, - help="start index", - ) - parser.add_argument( - "--end", - type=int, - default=-1, - help="end index", - ) - args = parser.parse_args() - if args.outdir is not None: os.makedirs(args.outdir, exist_ok=True) - # Initialize the ICGANInference class config = get_config() + config.seed = args.ith_sample icgan_inference = ICGANInference(config) - transform = transforms.Compose([ data_utils.CenterCropLongEdge(), - # transforms.Resize((args.img_save_size, args.img_save_size)), transforms.ToTensor(), transforms.Normalize(config.norm_mean, config.norm_std)]) @@ -88,21 +87,13 @@ def main(): assert args.start < n assert args.end <= n - for i in range(args.start, args.end): - print(f"Batch {i}") + for i in tqdm(range(args.start, args.end)): batch = imagenet_dataset[i] images = batch[0] - start = time.time() generated_images = icgan_inference.run_inference(input_image_tensor=images.unsqueeze(0)) - end = time.time() - print(f"Time taken to generate images: {end-start} seconds") - ## save images path = imagenet_dataset.samples[i][0] - start = time.time() save_images(path, generated_images, args.outdir) - end = time.time() - print(f"Time taken to save images: {end-start} seconds") if __name__ == "__main__": main() \ No newline at end of file