Update cldm/LumiNet.py

cldm/LumiNet.py

@@ -2,9 +2,6 @@ import einops
 import torch
 import torch as th
 import torch.nn as nn
-# from cldm.latent_intrinsic_new import LatentIntrinsc
-# from cldm.latent_intrinsic import LatentIntrinsc
-# from cldm.latent_intrinsic_crossattn import LatentIntrinsc
 from cldm.latent_intrinsic import LatentIntrinsc
 from ldm.modules.diffusionmodules.util import (
     conv_nd,
@@ -50,7 +47,6 @@ class ControlledUnetModel(UNetModel):
                 h = torch.cat([h, hs.pop() + control.pop()], dim=1)
             h = module(h, emb, context)
             # print("fool! that is the shape of the context! at output",context.shape)
-            # exit(0)
 
         h = h.type(x.dtype)
         return self.out(h)
@@ -173,60 +169,6 @@ class ControlNet(nn.Module):
             zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
         )
 
-        # self.input_latent_hint_block = TimestepEmbedSequential(
-        #     # conv_nd(dims, hint_channels, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 32, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 32, 32, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 32, 96, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 96, 96, 3, padding=1),
-        #     # nn.SiLU(),
-        #     conv_nd(dims, 128, 256, 3, padding=1, stride=1),
-        #     nn.SiLU(),
-        #     zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
-        # )
-
-        # self.input_latent_hint_cat_block = TimestepEmbedSequential(
-        #     # conv_nd(dims, hint_channels, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 32, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 144, 144, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 144, 256, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     conv_nd(dims, 144, 144, 3, padding=1),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 144, 256, 3, padding=1, stride=1),
-        #     nn.SiLU(),
-        #     zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
-        # )
-
-        # self.input_latent_hint_cat_block = TimestepEmbedSequential(
-        #     # conv_nd(dims, hint_channels, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 32, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     conv_nd(dims, 144, 144, 3, padding=1),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 144, 256, 3, padding=1, stride=2),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 256, 256, 3, padding=1),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 256, 256, 3, padding=1, stride=1),
-        #     nn.SiLU(),
-        #     zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
-        # )
-
         self.input_latent_hint_cat_atten_block = TimestepEmbedSequential(
             # conv_nd(dims, hint_channels, 16, 3, padding=1),
             # nn.SiLU(),
@@ -245,41 +187,6 @@ class ControlNet(nn.Module):
             zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
         )
 
-        # self.input_latent_hint_crossattn_block = TimestepEmbedSequential(
-        #     # conv_nd(dims, hint_channels, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 32, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     conv_nd(dims, 128, 128, 3, padding=1),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 128, 256, 3, padding=1, stride=2),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 256, 256, 3, padding=1),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 256, 256, 3, padding=1, stride=1),
-        #     nn.SiLU(),
-        #     zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
-        # )
-
-        # self.input_latent_hint_cat_eq_block = TimestepEmbedSequential(
-        #     # conv_nd(dims, hint_channels, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 16, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 16, 32, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 32, 32, 3, padding=1),
-        #     # nn.SiLU(),
-        #     # conv_nd(dims, 32, 96, 3, padding=1, stride=2),
-        #     # nn.SiLU(),
-        #     conv_nd(dims, 256, 512, 3, padding=1),
-        #     nn.SiLU(),
-        #     conv_nd(dims, 512, 256, 3, padding=1, stride=1),
-        #     nn.SiLU(),
-        #     zero_module(conv_nd(dims, 256, model_channels, 3, padding=1))
-        # )
 
         self._feature_size = model_channels
         input_block_chans = [model_channels]
@@ -398,15 +305,7 @@ class ControlNet(nn.Module):
         self._feature_size += ch
         self.latent_iid = True
         self.concat = True
-        # Old pure ControlNet
-        # if self.latent_iid:
-        #     # print(hint.shape)
-        #     if self.concat:
-        #         self.input_hint_block = self.input_latent_hint_cat_block
-        #     else:
-        #         self.input_hint_block = self.input_latent_hint_block
-
-        # New w. the crossattn version #comment before train Uncomment for test
+    
         if self.latent_iid:
             self.input_hint_block = self.input_latent_hint_cat_atten_block
     def make_zero_conv(self, channels):
@@ -414,31 +313,20 @@ class ControlNet(nn.Module):
 #our modification for the latent intrinsic
     def add_latent_prior(self):
         self.prior_extracter = LatentIntrinsc()
-        # send the extracted information to control net encoder (original image)
-        # self.input_hint_block[0] = conv_nd(2, 19, 16, 3, padding=1)
+    
 
     def forward(self, x, hint, timesteps, context, **kwargs):
         t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
         emb = self.time_embed(t_emb)
         self.latent_iid = True
-        # self.concat = False
         # adding a cross-attention version
         if self.latent_iid:
             hint, hint_lighting = self.prior_extracter(hint)
             self.new_context = torch.cat([context,hint_lighting],1)
-            # self.new_context = context*hint_lighting
             context = self.new_context 
-            # print("hint_context",context.shape)
-            # exit(0)
+
             self.input_hint_block = self.input_latent_hint_cat_atten_block
-        # pure controlnet with latent guidance
-        # if self.latent_iid:
-        #     hint = self.prior_extracter(hint)
-        #     # print(hint.shape)
-        #     if self.concat:
-        #         self.input_hint_block = self.input_latent_hint_cat_block
-        #     else:
-        #         self.input_hint_block = self.input_latent_hint_block
+    
             guided_hint = self.input_hint_block(hint, emb, context)
         else:
             guided_hint = self.input_hint_block(hint, emb, context)
@@ -470,6 +358,19 @@ class ControlLDM(LatentDiffusion):
         self.only_mid_control = only_mid_control
         self.control_scales = [1.0] * 13
 
+    # load bypass decoder
+    @torch.no_grad()
+    def change_first_stage(self, checkpoint_file,og=False):
+        del self.first_stage_model
+        from modi_vae.autoencoder import AutoencoderKL
+        self.first_stage_model = AutoencoderKL(load_checkpoint=False)
+        state_dict = torch.load(checkpoint_file, map_location=torch.device("cpu"))["state_dict"]
+        new_state_dict = {}
+        for s in state_dict:
+            new_state_dict[s]=state_dict[s]
+        
+        self.first_stage_model.load_state_dict(new_state_dict)
+        print("Successfully load new auto-encoder")
 
     @torch.no_grad()
     def add_new_layers(self):
@@ -491,8 +392,6 @@ class ControlLDM(LatentDiffusion):
         diffusion_model = self.model.diffusion_model
 
         cond_txt = torch.cat(cond['c_crossattn'], 1)
-        # print('attention_shape:',cond)
-        # exit(0)
         if cond['c_concat'] is None:
             eps = diffusion_model(x=x_noisy, timesteps=t, context=cond_txt, control=None, only_mid_control=self.only_mid_control)
         else:
@@ -504,13 +403,10 @@ class ControlLDM(LatentDiffusion):
             eps = diffusion_model(x=x_noisy, timesteps=t, context=cond_txt, control=control, only_mid_control=self.only_mid_control)
         # print('unet_context',cond_txt.shape)
         return eps
-    # @torch.no_grad()
-    # def get_attn_intrinsic(self):
-    #     return self.control_model.new_context
     @torch.no_grad()
     def get_unconditional_conditioning(self, N):
         return self.get_learned_conditioning([""] * N)
-
+    
     @torch.no_grad()
     def log_images(self, batch, N=4, n_row=2, sample=False, ddim_steps=50, ddim_eta=0.0, return_keys=None,
                    quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True,
@@ -583,9 +479,7 @@ class ControlLDM(LatentDiffusion):
     def configure_optimizers(self):
         lr = self.learning_rate
         params = list(self.control_model.parameters())
-        # if not self.sd_locked:
-        #     params += list(self.model.diffusion_model.output_blocks.parameters())
-        #     params += list(self.model.diffusion_model.out.parameters())
+    
         if self.crossattn_start: #here we also train the cross-attan in the input layer if has any
             for block in self.model.diffusion_model.input_blocks:
                 for layer in block: