remove ema parts in checkpoints.

app.py

@@ -66,11 +66,11 @@ def process_multi_wrapper_only_show_rendered(rendered_txt_0, rendered_txt_1, ren
                                      shared_eta, shared_a_prompt, shared_n_prompt, 
                                      only_show_rendered_image=True)  
 
-# cfg = OmegaConf.load("config.yaml")
-# model = load_model_from_config(cfg, "model_states.pt", verbose=True)
 
 cfg = OmegaConf.load("config.yaml")
-model = load_model_from_config(cfg, "model.ckpt", verbose=True)
+model = load_model_from_config(cfg, "model_wo_ema.ckpt", verbose=True)
+# model = load_model_from_config(cfg, "model_states.pt", verbose=True)
+# model = load_model_from_config(cfg, "model.ckpt", verbose=True)
 
 ddim_sampler = DDIMSampler(model)
 render_tool = Render_Text(model)

model_states.pt → model_wo_ema.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2b56f1251182afabc8d5291e07c3a3aaf21d85d36445b25d0057fc1960d63de5
-size 9880058178
+oid sha256:0b86b22188bf580e80773a5ae101bf9787eb258349f3f1acf0ae50fd10cb3fec
+size 6671922039

transfer.py

@@ -6,9 +6,15 @@ model = load_model_from_config(cfg, "model_states.pt", verbose=True)
 
 from pytorch_lightning.callbacks import ModelCheckpoint
 with model.ema_scope("store ema weights"):
+    model_sd = model.state_dict()
+    store_sd = {}
+    for key in model_sd:
+        if "ema" in key:
+            continue
+        store_sd[key] = model_sd[key]
     file_content = {
-        'state_dict': model.state_dict()
+        'state_dict': store_sd
     }
-    torch.save(file_content, "model.ckpt")
+    torch.save(file_content, "model_wo_ema.ckpt")
     print("has stored the transfered ckpt.")
 print("trial ends!")