Spaces:
Runtime error
Runtime error
| import os | |
| from glob import glob | |
| from subprocess import call | |
| import json | |
| def base(): | |
| return { | |
| "slurm":{ | |
| "t": 360, | |
| "N": 2, | |
| "n": 8, | |
| }, | |
| "model":{ | |
| "dataset" :"wds", | |
| "dataset_root": "/p/scratch/ccstdl/cherti1/CC12M/{00000..01099}.tar", | |
| "image_size": 256, | |
| "num_channels": 3, | |
| "num_channels_dae": 128, | |
| "ch_mult": "1 1 2 2 4 4", | |
| "num_timesteps": 4, | |
| "num_res_blocks": 2, | |
| "batch_size": 8, | |
| "num_epoch": 1000, | |
| "ngf": 64, | |
| "embedding_type": "positional", | |
| "use_ema": "", | |
| "ema_decay": 0.999, | |
| "r1_gamma": 1.0, | |
| "z_emb_dim": 256, | |
| "lr_d": 1e-4, | |
| "lr_g": 1.6e-4, | |
| "lazy_reg": 10, | |
| "save_content": "", | |
| "save_ckpt_every": 1, | |
| "masked_mean": "", | |
| "resume": "", | |
| }, | |
| } | |
| def ddgan_cc12m_v2(): | |
| cfg = base() | |
| cfg['slurm']['N'] = 2 | |
| cfg['slurm']['n'] = 8 | |
| return cfg | |
| def ddgan_cc12m_v6(): | |
| cfg = base() | |
| cfg['model']['text_encoder'] = "google/t5-v1_1-large" | |
| return cfg | |
| def ddgan_cc12m_v7(): | |
| cfg = base() | |
| cfg['model']['classifier_free_guidance_proba'] = 0.2 | |
| cfg['slurm']['N'] = 2 | |
| cfg['slurm']['n'] = 8 | |
| return cfg | |
| def ddgan_cc12m_v8(): | |
| cfg = base() | |
| cfg['model']['text_encoder'] = "google/t5-v1_1-large" | |
| cfg['model']['classifier_free_guidance_proba'] = 0.2 | |
| return cfg | |
| def ddgan_cc12m_v9(): | |
| cfg = base() | |
| cfg['model']['text_encoder'] = "google/t5-v1_1-large" | |
| cfg['model']['classifier_free_guidance_proba'] = 0.2 | |
| cfg['model']['num_channels_dae'] = 320 | |
| cfg['model']['image_size'] = 64 | |
| cfg['model']['batch_size'] = 1 | |
| return cfg | |
| def ddgan_cc12m_v11(): | |
| cfg = base() | |
| cfg['model']['text_encoder'] = "google/t5-v1_1-large" | |
| cfg['model']['classifier_free_guidance_proba'] = 0.2 | |
| cfg['model']['cross_attention'] = "" | |
| return cfg | |
| def ddgan_cc12m_v12(): | |
| cfg = ddgan_cc12m_v11() | |
| cfg['model']['text_encoder'] = "google/t5-v1_1-xl" | |
| cfg['model']['preprocessing'] = 'random_resized_crop_v1' | |
| return cfg | |
| def ddgan_cc12m_v13(): | |
| cfg = ddgan_cc12m_v12() | |
| cfg['model']['discr_type'] = "large_cond_attn" | |
| return cfg | |
| def ddgan_cc12m_v14(): | |
| cfg = ddgan_cc12m_v12() | |
| cfg['model']['num_channels_dae'] = 192 | |
| return cfg | |
| def ddgan_cc12m_v15(): | |
| cfg = ddgan_cc12m_v11() | |
| cfg['model']['mismatch_loss'] = '' | |
| cfg['model']['grad_penalty_cond'] = '' | |
| return cfg | |
| def ddgan_cifar10_cond17(): | |
| cfg = base() | |
| cfg['model']['image_size'] = 32 | |
| cfg['model']['classifier_free_guidance_proba'] = 0.2 | |
| cfg['model']['ch_mult'] = "1 2 2 2" | |
| cfg['model']['cross_attention'] = "" | |
| cfg['model']['dataset'] = "cifar10" | |
| cfg['model']['n_mlp'] = 4 | |
| return cfg | |
| def ddgan_cifar10_cond18(): | |
| cfg = ddgan_cifar10_cond17() | |
| cfg['model']['text_encoder'] = "google/t5-v1_1-xl" | |
| return cfg | |
| def ddgan_cifar10_cond19(): | |
| cfg = ddgan_cifar10_cond17() | |
| cfg['model']['discr_type'] = 'small_cond_attn' | |
| cfg['model']['mismatch_loss'] = '' | |
| cfg['model']['grad_penalty_cond'] = '' | |
| return cfg | |
| def ddgan_laion_aesthetic_v1(): | |
| cfg = ddgan_cc12m_v11() | |
| cfg['model']['dataset_root'] = '"/p/scratch/ccstdl/cherti1/LAION-aesthetic/output/{00000..05038}.tar"' | |
| return cfg | |
| def ddgan_laion_aesthetic_v2(): | |
| cfg = ddgan_laion_aesthetic_v1() | |
| cfg['model']['discr_type'] = "large_cond_attn" | |
| return cfg | |
| def ddgan_laion_aesthetic_v3(): | |
| cfg = ddgan_laion_aesthetic_v1() | |
| cfg['model']['text_encoder'] = "google/t5-v1_1-xl" | |
| cfg['model']['mismatch_loss'] = '' | |
| cfg['model']['grad_penalty_cond'] = '' | |
| return cfg | |
| def ddgan_laion_aesthetic_v4(): | |
| cfg = ddgan_laion_aesthetic_v1() | |
| cfg['model']['text_encoder'] = "openclip/ViT-L-14-336/openai" | |
| return cfg | |
| def ddgan_laion_aesthetic_v5(): | |
| cfg = ddgan_laion_aesthetic_v1() | |
| cfg['model']['mismatch_loss'] = '' | |
| cfg['model']['grad_penalty_cond'] = '' | |
| return cfg | |
| def ddgan_laion2b_v1(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['mismatch_loss'] = '' | |
| cfg['model']['grad_penalty_cond'] = '' | |
| cfg['model']['num_channels_dae'] = 224 | |
| cfg['model']['batch_size'] = 2 | |
| cfg['model']['discr_type'] = "large_cond_attn" | |
| cfg['model']['preprocessing'] = 'random_resized_crop_v1' | |
| return cfg | |
| def ddgan_laion_aesthetic_v6(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['no_lr_decay'] = '' | |
| return cfg | |
| def ddgan_laion_aesthetic_v7(): | |
| cfg = ddgan_laion_aesthetic_v6() | |
| cfg['model']['r1_gamma'] = 5 | |
| return cfg | |
| def ddgan_laion_aesthetic_v8(): | |
| cfg = ddgan_laion_aesthetic_v6() | |
| cfg['model']['num_timesteps'] = 8 | |
| return cfg | |
| def ddgan_laion_aesthetic_v9(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['num_channels_dae'] = 384 | |
| return cfg | |
| def ddgan_sd_v1(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| return cfg | |
| def ddgan_sd_v2(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| return cfg | |
| def ddgan_sd_v3(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| return cfg | |
| def ddgan_sd_v4(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| return cfg | |
| def ddgan_sd_v5(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['num_timesteps'] = 8 | |
| return cfg | |
| def ddgan_sd_v6(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['num_channels_dae'] = 192 | |
| return cfg | |
| def ddgan_sd_v7(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| return cfg | |
| def ddgan_sd_v8(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['image_size'] = 512 | |
| return cfg | |
| def ddgan_laion_aesthetic_v12(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| return cfg | |
| def ddgan_laion_aesthetic_v13(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['text_encoder'] = "openclip/ViT-H-14/laion2b_s32b_b79k" | |
| return cfg | |
| def ddgan_laion_aesthetic_v14(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['text_encoder'] = "openclip/ViT-H-14/laion2b_s32b_b79k" | |
| return cfg | |
| def ddgan_sd_v9(): | |
| cfg = ddgan_laion_aesthetic_v3() | |
| cfg['model']['text_encoder'] = "openclip/ViT-H-14/laion2b_s32b_b79k" | |
| return cfg | |
| def ddgan_sd_v10(): | |
| cfg = ddgan_sd_v9() | |
| cfg['model']['num_timesteps'] = 2 | |
| return cfg | |
| def ddgan_laion2b_v2(): | |
| cfg = ddgan_sd_v9() | |
| return cfg | |
| def ddgan_ddb_v1(): | |
| cfg = ddgan_sd_v10() | |
| return cfg | |
| def ddgan_sd_v11(): | |
| cfg = ddgan_sd_v10() | |
| cfg['model']['image_size'] = 512 | |
| return cfg | |
| def ddgan_ddb_v2(): | |
| cfg = ddgan_ddb_v1() | |
| cfg['model']['num_timesteps'] = 1 | |
| return cfg | |
| def ddgan_ddb_v3(): | |
| cfg = ddgan_ddb_v1() | |
| cfg['model']['num_channels_dae'] = 192 | |
| cfg['model']['num_timesteps'] = 2 | |
| return cfg | |
| models = [ | |
| ddgan_cifar10_cond17, # cifar10, cross attn for discr | |
| ddgan_cifar10_cond18, # cifar10, xl encoder | |
| ddgan_cifar10_cond19, # cifar10, xl encoder | |
| ddgan_cc12m_v2, # baseline (no large text encoder, no classifier guidance) | |
| ddgan_cc12m_v6, # like v2 but using large T5 text encoder | |
| ddgan_cc12m_v7, # like v2 but with classifier guidance | |
| ddgan_cc12m_v8, # like v6 but classifier guidance | |
| ddgan_cc12m_v9, # ~1B model but 64x64 resolution | |
| ddgan_cc12m_v11, # large text encoder + cross attention + classifier free guidance | |
| ddgan_cc12m_v12, # T5-XL + cross attention + classifier free guidance + random_resized_crop_v1 | |
| ddgan_cc12m_v13, # T5-XL + cross attention + classifier free guidance + random_resized_crop_v1 + cond attn | |
| ddgan_cc12m_v14, # T5-XL + cross attention + classifier free guidance + random_resized_crop_v1 + 300M model | |
| ddgan_cc12m_v15, # fine-tune v11 with --mismatch_loss and --grad_penalty_cond | |
| ddgan_laion_aesthetic_v1, # like ddgan_cc12m_v11 but fine-tuned on laion aesthetic | |
| ddgan_laion_aesthetic_v2, # like ddgan_laion_aesthetic_v1 but trained from scratch with the new cross attn discr | |
| ddgan_laion_aesthetic_v3, # like ddgan_laion_aesthetic_v1 but trained from scratch with T5-XL (continue from 23aug with mismatch and grad penalty and random_resized_crop_v1) | |
| ddgan_laion_aesthetic_v4, # like ddgan_laion_aesthetic_v1 but trained from scratch with OpenAI's ClipEncoder | |
| ddgan_laion_aesthetic_v5, # fine-tune ddgan_laion_aesthetic_v1 with mismatch and cond grad penalty losses | |
| ddgan_laion_aesthetic_v6, # like v3 but without lr decay | |
| ddgan_laion_aesthetic_v7, # like v6 but with r1 gamma of 5 instead of 1, trying to constrain the discr more. | |
| ddgan_laion_aesthetic_v8, # like v6 but with 8 timesteps | |
| ddgan_laion_aesthetic_v9, | |
| ddgan_laion_aesthetic_v12, | |
| ddgan_laion_aesthetic_v13, | |
| ddgan_laion_aesthetic_v14, | |
| ddgan_laion2b_v1, | |
| ddgan_sd_v1, | |
| ddgan_sd_v2, | |
| ddgan_sd_v3, | |
| ddgan_sd_v4, | |
| ddgan_sd_v5, | |
| ddgan_sd_v6, | |
| ddgan_sd_v7, | |
| ddgan_sd_v8, | |
| ddgan_sd_v9, | |
| ddgan_sd_v10, | |
| ddgan_sd_v11, | |
| ddgan_laion2b_v2, | |
| ddgan_ddb_v1, | |
| ddgan_ddb_v2, | |
| ddgan_ddb_v3 | |
| ] | |
| def get_model(model_name): | |
| for model in models: | |
| if model.__name__ == model_name: | |
| return model() | |
| def test(model_name, *, cond_text="", batch_size:int=None, epoch:int=None, guidance_scale:float=0, fid=False, real_img_dir="", q=0.0, seed=0, nb_images_for_fid=0, scale_factor_h=1, scale_factor_w=1, compute_clip_score=False, eval_name="", scale_method="convolutional"): | |
| cfg = get_model(model_name) | |
| model = cfg['model'] | |
| if epoch is None: | |
| paths = glob('./saved_info/dd_gan/{}/{}/netG_*.pth'.format(model["dataset"], model_name)) | |
| epoch = max( | |
| [int(os.path.basename(path).replace(".pth", "").split("_")[1]) for path in paths] | |
| ) | |
| args = {} | |
| args['exp'] = model_name | |
| args['image_size'] = model['image_size'] | |
| args['seed'] = seed | |
| args['num_channels'] = model['num_channels'] | |
| args['dataset'] = model['dataset'] | |
| args['num_channels_dae'] = model['num_channels_dae'] | |
| args['ch_mult'] = model['ch_mult'] | |
| args['num_timesteps'] = model['num_timesteps'] | |
| args['num_res_blocks'] = model['num_res_blocks'] | |
| args['batch_size'] = model['batch_size'] if batch_size is None else batch_size | |
| args['epoch'] = epoch | |
| args['cond_text'] = f'"{cond_text}"' | |
| args['text_encoder'] = model.get("text_encoder") | |
| args['cross_attention'] = model.get("cross_attention") | |
| args['guidance_scale'] = guidance_scale | |
| args['masked_mean'] = model.get("masked_mean") | |
| args['dynamic_thresholding_quantile'] = q | |
| args['scale_factor_h'] = scale_factor_h | |
| args['scale_factor_w'] = scale_factor_w | |
| args['n_mlp'] = model.get("n_mlp") | |
| args['scale_method'] = scale_method | |
| if fid: | |
| args['compute_fid'] = '' | |
| args['real_img_dir'] = real_img_dir | |
| args['nb_images_for_fid'] = nb_images_for_fid | |
| if compute_clip_score: | |
| args['compute_clip_score'] = "" | |
| if eval_name: | |
| args["eval_name"] = eval_name | |
| cmd = "python -u test_ddgan.py " + " ".join(f"--{k} {v}" for k, v in args.items() if v is not None) | |
| print(cmd) | |
| call(cmd, shell=True) | |
| def eval_results(model_name): | |
| import pandas as pd | |
| rows = [] | |
| cfg = get_model(model_name) | |
| model = cfg['model'] | |
| paths = glob('./saved_info/dd_gan/{}/{}/fid*.json'.format(model["dataset"], model_name)) | |
| for path in paths: | |
| with open(path, "r") as fd: | |
| data = json.load(fd) | |
| row = {} | |
| row['fid'] = data['fid'] | |
| row['epoch'] = data['epoch_id'] | |
| rows.append(row) | |
| out = './saved_info/dd_gan/{}/{}/fid.csv'.format(model["dataset"], model_name) | |
| df = pd.DataFrame(rows) | |
| df.to_csv(out, index=False) | |
| if __name__ == "__main__": | |
| from clize import run | |
| run([test, eval_results]) | |