Create cool_models.py

cool_models.py

@@ -0,0 +1,123 @@
+import torch
+from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults
+import lpips
+import clip
+
+
+from encoders.modules import BERTEmbedder
+from models.clipseg import CLIPDensePredT
+
+from huggingface_hub import hf_hub_download
+
+STEPS = 100
+USE_DDPM = False
+USE_DDIM = False
+USE_CPU = False
+CLIP_SEG_PATH = './weights/rd64-uni.pth'
+CLIP_GUIDANCE = False
+
+def make_models():
+    segmodel = CLIPDensePredT(version='ViT-B/16', reduce_dim=64)
+    segmodel.eval()
+
+    # non-strict, because we only stored decoder weights (not CLIP weights)
+    segmodel.load_state_dict(torch.load(CLIP_SEG_PATH, map_location=torch.device('cpu')), strict=False)
+    # segmodel.save_pretrained("./weights/hf_clipseg")
+
+    device = torch.device('cuda:0' if (torch.cuda.is_available() and not USE_CPU) else 'cpu')
+    print('Using device:', device)
+
+    hf_inpaint_path = hf_hub_download("alvanlii/rdm_inpaint", "inpaint.pt")
+    model_state_dict = torch.load(hf_inpaint_path, map_location='cpu')
+
+    # print(
+    #     'hey',
+    #     'clip_proj.weight' in model_state_dict, # True
+    #     model_state_dict['input_blocks.0.0.weight'].shape[1] == 8, # True
+    #     'external_block.0.0.weight' in model_state_dict # False
+    # )
+
+    model_params = {
+        'attention_resolutions': '32,16,8',
+        'class_cond': False,
+        'diffusion_steps': 1000,
+        'rescale_timesteps': True,
+        'timestep_respacing': STEPS,  # Modify this value to decrease the number of
+                                    # timesteps.
+        'image_size': 32,
+        'learn_sigma': False,
+        'noise_schedule': 'linear',
+        'num_channels': 320,
+        'num_heads': 8,
+        'num_res_blocks': 2,
+        'resblock_updown': False,
+        'use_fp16': False,
+        'use_scale_shift_norm': False,
+        'clip_embed_dim': 768,
+        'image_condition': True,
+        'super_res_condition': False,
+    }
+
+    if USE_DDPM:
+        model_params['timestep_respacing'] = '1000'
+    if USE_DDIM:
+        if STEPS:
+            model_params['timestep_respacing'] = 'ddim'+str(STEPS)
+        else:
+            model_params['timestep_respacing'] = 'ddim50'
+    elif STEPS:
+        model_params['timestep_respacing'] = str(STEPS)
+
+    model_config = model_and_diffusion_defaults()
+    model_config.update(model_params)
+
+    if USE_CPU:
+        model_config['use_fp16'] = False
+
+
+    model, diffusion = create_model_and_diffusion(**model_config)
+    model.load_state_dict(model_state_dict, strict=False)
+
+    model.requires_grad_(CLIP_GUIDANCE).eval().to(device)
+
+    if model_config['use_fp16']:
+        model.convert_to_fp16()
+    else:
+        model.convert_to_fp32()
+
+    def set_requires_grad(model, value):
+        for param in model.parameters():
+            param.requires_grad = value
+
+
+    lpips_model = lpips.LPIPS(net="vgg").to(device)
+    hf_kl_path = hf_hub_download("alvanlii/rdm_inpaint", "kl-f8.pt")
+
+    ldm = torch.load(hf_kl_path, map_location="cpu")
+    
+    # torch.save(ldm, "./weights/hf_ldm")
+    ldm.to(device)
+    ldm.eval()
+    ldm.requires_grad_(CLIP_GUIDANCE)
+    set_requires_grad(ldm, CLIP_GUIDANCE)
+
+    bert = BERTEmbedder(1280, 32)
+    hf_bert_path = hf_hub_download("alvanlii/rdm_inpaint", 'bert.pt')
+    # bert = BERTEmbedder.from_pretrained("alvanlii/rdm_bert")
+    sd = torch.load(hf_bert_path, map_location="cpu")
+    bert.load_state_dict(sd)
+    # bert.save_pretrained("./weights/hf_bert")
+
+    bert.to(device)
+    bert.half().eval()
+    set_requires_grad(bert, False)
+
+
+    clip_model, clip_preprocess = clip.load('ViT-L/14', device=device, jit=False)
+    clip_model.eval().requires_grad_(False)
+
+    return segmodel, model, diffusion, ldm, bert, clip_model, model_params
+
+
+if __name__ == "__main__":
+    make_models()