csyxwei commited on
Commit
e1cbfc8
1 Parent(s): eae8580

fix the OOM problem

Browse files
Files changed (1) hide show
  1. model.py +3 -1
model.py CHANGED
@@ -29,7 +29,7 @@ snapshot_download('ELITE-library/ELITE',
29
  sys.path.insert(0, submodule_dir.as_posix())
30
 
31
  from train_local import (Mapper, MapperLocal, inj_forward_crossattention,
32
- inj_forward_text, th2image)
33
 
34
 
35
  def get_tensor_clip(normalize=True, toTensor=True):
@@ -335,6 +335,7 @@ class Model:
335
  placeholder_idx.detach(),
336
  'LAMBDA': lambda_
337
  }).sample
 
338
  latent_model_input = self.scheduler.scale_model_input(latents, t)
339
 
340
  noise_pred_uncond = self.unet(latent_model_input,
@@ -343,6 +344,7 @@ class Model:
343
  'CONTEXT_TENSOR':
344
  uncond_embeddings,
345
  }).sample
 
346
  noise_pred = noise_pred_uncond + guidance_scale * (
347
  noise_pred_text - noise_pred_uncond)
348
 
 
29
  sys.path.insert(0, submodule_dir.as_posix())
30
 
31
  from train_local import (Mapper, MapperLocal, inj_forward_crossattention,
32
+ inj_forward_text, th2image, value_local_list)
33
 
34
 
35
  def get_tensor_clip(normalize=True, toTensor=True):
 
335
  placeholder_idx.detach(),
336
  'LAMBDA': lambda_
337
  }).sample
338
+ value_local_list.clear()
339
  latent_model_input = self.scheduler.scale_model_input(latents, t)
340
 
341
  noise_pred_uncond = self.unet(latent_model_input,
 
344
  'CONTEXT_TENSOR':
345
  uncond_embeddings,
346
  }).sample
347
+ value_local_list.clear()
348
  noise_pred = noise_pred_uncond + guidance_scale * (
349
  noise_pred_text - noise_pred_uncond)
350