Update models/unet.py

models/unet.py

@@ -610,112 +610,10 @@ class UNet3DConditionModel(ModelMixin, ConfigMixin):
         # config["num_class_embeds"] = 100
 
         from diffusers.utils import WEIGHTS_NAME # diffusion_pytorch_model.bin
-        
-        # {'_class_name': 'UNet3DConditionModel', 
-        #  '_diffusers_version': '0.2.2', 
-        #  'act_fn': 'silu', 
-        #  'attention_head_dim': 8, 
-        #  'block_out_channels': [320, 640, 1280, 1280], 
-        #  'center_input_sample': False, 
-        #  'cross_attention_dim': 768, 
-        #  'down_block_types': 
-        #  ['CrossAttnDownBlock3D', 
-        #   'CrossAttnDownBlock3D', 
-        #   'CrossAttnDownBlock3D', 
-        #   'DownBlock3D'], 
-        #   'downsample_padding': 1, 
-        #   'flip_sin_to_cos': True, 
-        #   'freq_shift': 0, 
-        #   'in_channels': 4, 
-        #   'layers_per_block': 2, 
-        #   'mid_block_scale_factor': 1, 
-        #   'norm_eps': 1e-05, 
-        #   'norm_num_groups': 32, 
-        #   'out_channels': 4, 
-        #   'sample_size': 64, 
-        #   'up_block_types': 
-        #   ['UpBlock3D', 
-        #    'CrossAttnUpBlock3D', 
-        #    'CrossAttnUpBlock3D', 
-        #    'CrossAttnUpBlock3D']}
+
 
         model = cls.from_config(config)
-        # model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
-        # if not os.path.isfile(model_file):
-        #     raise RuntimeError(f"{model_file} does not exist")
-        # state_dict = torch.load(model_file, map_location="cpu")
-
-        # if use_concat:
-        #     new_state_dict = {}
-        #     conv_in_weight = state_dict["conv_in.weight"]
-        #     new_conv_weight = torch.zeros((conv_in_weight.shape[0], 9, *conv_in_weight.shape[2:]), dtype=conv_in_weight.dtype)
-        
-        #     for i, j in zip([0, 1, 2, 3], [0, 1, 2, 3, 4, 5, 6, 7, 8]):
-        #         new_conv_weight[:, j] = conv_in_weight[:, i]
-        #     new_state_dict["conv_in.weight"] = new_conv_weight
-        #     new_state_dict["conv_in.bias"] = state_dict["conv_in.bias"]
-        #     for k, v in model.state_dict().items():
-        #         # print(k)
-        #         if '_temp.' in k:
-        #             new_state_dict.update({k: v})
-        #         if 'attn_fcross' in k: # conpy parms of attn1 to attn_fcross
-        #             k = k.replace('attn_fcross', 'attn1')
-        #             state_dict.update({k: state_dict[k]})
-        #         if 'norm_fcross' in k:
-        #             k = k.replace('norm_fcross', 'norm1')
-        #             state_dict.update({k: state_dict[k]})
-                    
-        #         if 'conv_in' in k:
-        #             continue
-        #         else:
-        #             new_state_dict[k] = v
-        #         # # tmp 
-        #         # if 'class_embedding' in k:
-        #         #     state_dict.update({k: v})
-        #     # breakpoint()
-        #     model.load_state_dict(new_state_dict)
-        # else:
-        #     for k, v in model.state_dict().items():
-        #         # print(k)
-        #         if '_temp' in k:
-        #             state_dict.update({k: v})
-        #         if 'attn_fcross' in k: # conpy parms of attn1 to attn_fcross
-        #             k = k.replace('attn_fcross', 'attn1')
-        #             state_dict.update({k: state_dict[k]})
-        #         if 'norm_fcross' in k:
-        #             k = k.replace('norm_fcross', 'norm1')
-        #             state_dict.update({k: state_dict[k]})
-
-        #     model.load_state_dict(state_dict)
 
-        return model
-    
-if __name__ == '__main__':
-    import torch
-    # from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # pretrained_model_path = "/mnt/petrelfs/maxin/work/pretrained/stable-diffusion-2-1-base/" # p cluster
-    pretrained_model_path = "/mnt/petrelfs/share_data/zhanglingjun/stable-diffusion-v1-4/" # p cluster
-    unet = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet").to(device)
-    # unet.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
-    unet.enable_xformers_memory_efficient_attention()
-    unet.enable_gradient_checkpointing()
-
-    unet.train()
-
-    use_image_num = 5
-    noisy_latents = torch.randn((2, 4, 16 + use_image_num, 32, 32)).to(device)
-    bsz = noisy_latents.shape[0]
-    timesteps = torch.randint(0, 1000, (bsz,)).to(device)
-    timesteps = timesteps.long()
-    encoder_hidden_states = torch.randn((bsz, 1 + use_image_num, 77, 768)).to(device)
-    # class_labels = torch.randn((bsz, )).to(device)
-    
 
-    model_pred = unet(sample=noisy_latents, timestep=timesteps, 
-                      encoder_hidden_states=encoder_hidden_states, 
-                      class_labels=None,
-                      use_image_num=use_image_num).sample
-    print(model_pred.shape)
+        return model
+