BestWishYsh commited on
Commit
4d65360
1 Parent(s): 0739b68

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -0
app.py CHANGED
@@ -124,10 +124,16 @@ class MagicTimeController:
124
  converted_vae_checkpoint = convert_ldm_vae_checkpoint(dreambooth_state_dict, self.vae.config)
125
  self.vae.load_state_dict(converted_vae_checkpoint)
126
 
 
 
 
127
  converted_unet_checkpoint = convert_ldm_unet_checkpoint(dreambooth_state_dict, self.unet_model.config)
128
  self.unet = copy.deepcopy(self.unet_model)
129
  self.unet.load_state_dict(converted_unet_checkpoint, strict=False)
130
 
 
 
 
131
  text_model = copy.deepcopy(self.text_model)
132
  self.text_encoder = convert_ldm_clip_text_model(text_model, dreambooth_state_dict)
133
 
 
124
  converted_vae_checkpoint = convert_ldm_vae_checkpoint(dreambooth_state_dict, self.vae.config)
125
  self.vae.load_state_dict(converted_vae_checkpoint)
126
 
127
+ if self.unet is not None:
128
+ del self.unet # 删除旧模型的引用
129
+ torch.cuda.empty_cache()
130
  converted_unet_checkpoint = convert_ldm_unet_checkpoint(dreambooth_state_dict, self.unet_model.config)
131
  self.unet = copy.deepcopy(self.unet_model)
132
  self.unet.load_state_dict(converted_unet_checkpoint, strict=False)
133
 
134
+ if self.text_encoder is not None:
135
+ del self.text_encoder
136
+ torch.cuda.empty_cache()
137
  text_model = copy.deepcopy(self.text_model)
138
  self.text_encoder = convert_ldm_clip_text_model(text_model, dreambooth_state_dict)
139