update the checkpoints fine-tuned on TextCaps 5K.

app.py

@@ -82,12 +82,18 @@ def load_ckpt(model_ckpt = "LAION-Glyph-10M-Epoch-5"):
         time.sleep(2)
         print("empty the cuda cache")
 
-    if model_ckpt == "LAION-Glyph-1M":
-        model = load_model_ckpt(model, "laion1M_model_wo_ema.ckpt")
-    elif model_ckpt == "LAION-Glyph-10M-Epoch-5":
+    # if model_ckpt == "LAION-Glyph-1M":
+    #     model = load_model_ckpt(model, "laion1M_model_wo_ema.ckpt")
+    if model_ckpt == "LAION-Glyph-10M-Epoch-5":
         model = load_model_ckpt(model, "laion10M_epoch_5_model_wo_ema.ckpt")
     elif model_ckpt == "LAION-Glyph-10M-Epoch-6":
         model = load_model_ckpt(model, "laion10M_epoch_6_model_wo_ema.ckpt")
+    elif model_ckpt == "TextCaps-5K-Epoch-10":
+        model = load_model_ckpt(model, "textcaps5K_epoch_10_model_wo_ema.ckpt")
+    elif model_ckpt == "TextCaps-5K-Epoch-20":
+        model = load_model_ckpt(model, "textcaps5K_epoch_20_model_wo_ema.ckpt")
+    elif model_ckpt == "TextCaps-5K-Epoch-40":
+        model = load_model_ckpt(model, "textcaps5K_epoch_40_model_wo_ema.ckpt")
 
     render_tool = Render_Text(model)
     output_str = f"already change the model checkpoint to {model_ckpt}"
@@ -126,7 +132,7 @@ with block:
         only_show_rendered_image = gr.Number(value=1, visible=False)
     default_width = [0.3, 0.3, 0.3, 0.3]
     default_top_left_x = [0.35, 0.15, 0.15, 0.5]
-    default_top_left_y = [0.5, 0.25, 0.75, 0.75]
+    default_top_left_y = [0.4, 0.15, 0.65, 0.65]
     with gr.Column():  
             
         with gr.Row(): 
@@ -154,7 +160,8 @@ with block:
                 with gr.Accordion("Model Options", open=False):
                     with gr.Row():
                         # model_ckpt = gr.inputs.Dropdown(["LAION-Glyph-10M", "Textcaps5K-10"], label="Checkpoint", default = "LAION-Glyph-10M")
-                        model_ckpt = gr.inputs.Dropdown(["LAION-Glyph-10M-Epoch-6", "LAION-Glyph-10M-Epoch-5", "LAION-Glyph-1M"], label="Checkpoint", default = "LAION-Glyph-10M-Epoch-6")
+                        # model_ckpt = gr.inputs.Dropdown(["LAION-Glyph-10M-Epoch-6", "LAION-Glyph-10M-Epoch-5", "LAION-Glyph-1M"], label="Checkpoint", default = "LAION-Glyph-10M-Epoch-6")
+                        model_ckpt = gr.inputs.Dropdown(["LAION-Glyph-10M-Epoch-6", "LAION-Glyph-10M-Epoch-5", "TextCaps-5K-Epoch-10", "TextCaps-5K-Epoch-20", "TextCaps-5K-Epoch-40"], label="Checkpoint", default = "LAION-Glyph-10M-Epoch-6")
                         # load_button = gr.Button(value = "Load Checkpoint")
             
             with gr.Accordion("Shared Advanced Options", open=False):  

cldm/cldm.py

@@ -532,13 +532,6 @@ class ControlLDM(LatentDiffusion):
         self.freeze_glyph_image_encoder = model.freeze_image_encoder #image_encoder.freeze_model
         self.glyph_control_model = model
         self.glyph_image_encoder_type = model.image_encoder_type
-            # self.glyph_control_optim = torch.optim.AdamW([
-            #     {"params": gain_or_bias_params, "weight_decay": 0.}, # "lr": self.glycon_lr},
-            #     {"params": rest_params, "weight_decay": self.glycon_wd}  #, "lr": self.glycon_lr},
-            # ],
-            # lr = self.glycon_lr
-            # )
-            # params += list(model.image_encoder.parameters())
         
         
 
@@ -738,16 +731,6 @@ class ControlLDM(LatentDiffusion):
         if decoder_params is not None:
             params_wlr.append({"params": decoder_params, "lr": self.decoder_lr})
         
-        
-        # if not self.sep_lr:
-        #     opt = torch.optim.AdamW(params, lr=lr)
-        # else:
-        #     opt = torch.optim.AdamW(
-        #         [
-        #         {"params": params},
-        #         {"params": decoder_params, "lr": self.decoder_lr} 
-        #         ], lr=lr
-        #     )
         if not self.freeze_glyph_image_encoder:
             if self.glyph_image_encoder_type == "CLIP":
                 # assert self.sep_lr
@@ -866,20 +849,6 @@ class ControlLDM(LatentDiffusion):
                     if p.requires_grad and p.grad is not None:
                         grad_norm_v = p.grad.cpu().detach().norm().item()
                         gradnorm_list.append(grad_norm_v)
-            # for name, p in self.named_parameters():
-            #     if p.requires_grad and p.grad is not None:
-            #         grad_norm_v = p.grad.detach().norm().item()
-            #         gradnorm_list.append(grad_norm_v)
-            #         if "textemb_merge_model" in name:
-            #             self.log("all_gradients/{}_norm".format(name),
-            #             gradnorm_list[-1],
-            #             prog_bar=False, logger=True, on_step=True, on_epoch=False
-            #             )
-            #         # if grad_norm_v > 0.1:
-            #         #     print("the norm of gradient w.r.t {} > 0.1: {:.2f}".format
-            #         #           (
-            #         #         name, grad_norm_v
-            #         #           ))
             if len(gradnorm_list):
                 self.log("all_gradients/grad_norm_mean", 
                     np.mean(gradnorm_list), 
@@ -943,19 +912,4 @@ class ControlLDM(LatentDiffusion):
                         prog_bar=False, logger=True, on_step=True, on_epoch=False
                     )
             del gradnorm_list
-            del zeroconvs
-
-    # def freeze_unet(self):
-    #     # Have some bugs
-    #     self.model.eval()
-    #     # self.model.train = disabled_train
-    #     for param in self.model.parameters():
-    #         param.requires_grad = False
-
-    #     if not self.sd_locked:
-    #         self.model.diffusion_model.output_blocks.train()
-    #         self.model.diffusion_model.out.train()
-    #         for param in self.model.diffusion_model.out.parameters():
-    #             param.requires_grad = True
-    #         for param in self.model.diffusion_model.output_blocks.parameters():
-    #             param.requires_grad = True
+            del zeroconvs

config_ema_unlock.yaml

@@ -0,0 +1,88 @@
+model:
+  base_learning_rate: 1.0e-6 #1.0e-5 #1.0e-4
+  target: cldm.cldm.ControlLDM
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    control_key: "hint"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false
+    conditioning_key: crossattn
+    monitor: #val/loss_simple_ema
+    scale_factor: 0.18215
+    only_mid_control: False
+    sd_locked: False #True
+    use_ema: True #TODO: specify
+
+    control_stage_config:
+      target: cldm.cldm.ControlNet
+      params:
+        use_checkpoint: True
+        image_size: 32 # unused
+        in_channels: 4
+        hint_channels: 3
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_head_channels: 64 # need to fix for flash-attn
+        use_spatial_transformer: True
+        use_linear_in_transformer: True
+        transformer_depth: 1
+        context_dim: 1024
+        legacy: False
+
+    unet_config:
+      target: cldm.cldm.ControlledUnetModel
+      params:
+        use_checkpoint: True
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_head_channels: 64 # need to fix for flash-attn
+        use_spatial_transformer: True
+        use_linear_in_transformer: True
+        transformer_depth: 1
+        context_dim: 1024
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          #attn_type: "vanilla-xformers"
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+      params:
+        freeze: True
+        layer: "penultimate"
+        # device: "cpu" #TODO: specify

laion1M_model_wo_ema.ckpt → textcaps5K_epoch_10_model_wo_ema.ckpt

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

textcaps5K_epoch_20_model_wo_ema.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c887aa42db7afbed071629bcf5a07cfccdcdc80216475d8a2536fed75cc600
+size 6671914001

textcaps5K_epoch_40_model_wo_ema.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:511be806f6e44f9c33af75df181adacfb3a0bb71aac8df8b303fff36e8e97dae
+size 6671914001

transfer.py

@@ -2,10 +2,13 @@ from omegaconf import OmegaConf
 from scripts.rendertext_tool import Render_Text, load_model_from_config
 import torch
 
-cfg = OmegaConf.load("config_ema.yaml")
-# model = load_model_from_config(cfg, "model_states.pt", verbose=True)
-model = load_model_from_config(cfg, "mp_rank_00_model_states.pt", verbose=True)
+# cfg = OmegaConf.load("config_ema.yaml")
+# # model = load_model_from_config(cfg, "model_states.pt", verbose=True)
+# model = load_model_from_config(cfg, "mp_rank_00_model_states.pt", verbose=True)
 
+cfg = OmegaConf.load("config_ema_unlock.yaml")
+epoch_idx = 39
+model = load_model_from_config(cfg, "epoch={:0>6d}.ckpt".format(epoch_idx), verbose=True)
 
 from pytorch_lightning.callbacks import ModelCheckpoint
 with model.ema_scope("store ema weights"):
@@ -18,6 +21,6 @@ with model.ema_scope("store ema weights"):
     file_content = {
         'state_dict': store_sd
     }
-    torch.save(file_content, "model_wo_ema.ckpt")
+    torch.save(file_content, f"textcaps5K_epoch_{epoch_idx+1}_model_wo_ema.ckpt")
     print("has stored the transfered ckpt.")
 print("trial ends!")