Update gligen/ldm/models/diffusion/plms.py

gligen/ldm/models/diffusion/plms.py

@@ -1,6 +1,6 @@
-import os
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "0"
-os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+# import os
+# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "0"
+# os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
 import torch
 import numpy as np
 from tqdm import tqdm
@@ -192,13 +192,19 @@ class PLMSSampler(object):
         
         print("optimize", index1)
         self.model.train()
-        torch.cuda.empty_cache()
+        # torch.cuda.empty_cache()
         while loss.item() > loss_threshold and iteration < max_iter and (index1 < max_index) :
             print('iter', iteration)
             # import pdb; pdb.set_trace()
             x = x.requires_grad_(True)
             input['x'] = x
             e_t,  att_first, att_second, att_third, self_first, self_second, self_third = self.model(input)
+
+            del self_first
+            del self_second
+            del self_third
+            del e_t
+            
             bboxes = input['boxes_att']
             object_positions = input['object_position'] 
             # loss1 = caculate_loss_self_att(self_first, self_second, self_third, bboxes=bboxes,
@@ -213,9 +219,7 @@ class PLMSSampler(object):
             del att_first
             del att_second
             del att_third
-            del self_first
-            del self_second
-            del self_third
+
             
             # grad_cond = x.grad 
             x = x - grad_cond