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

gligen/ldm/models/diffusion/plms.py

@@ -119,55 +119,95 @@ class PLMSSampler(object):
         return img
 
 
-    def update_loss_LoCo(self, input,index1, index, ts, type_loss='self_accross'):
+    # def update_loss_LoCo(self, input,index1, index, ts, type_loss='self_accross'):
         
-        # loss_scale = 30
-        # max_iter = 5
-        #print('time_factor is: ', time_factor)
+    #     # loss_scale = 30
+    #     # max_iter = 5
+    #     #print('time_factor is: ', time_factor)
+    #     if index1 < 10:
+    #         loss_scale = 8
+    #         max_iter = 5
+    #     elif index1 < 20:
+    #         loss_scale = 5
+    #         max_iter = 5
+    #     else:
+    #         loss_scale = 1
+    #         max_iter = 1
+    #     loss_threshold = 0.1
+        
+    #     max_index = 30
+    #     x = deepcopy(input["x"]) 
+    #     iteration = 0
+    #     loss = torch.tensor(10000)
+    #     input["timesteps"] = ts
+        
+    #     # print("optimize", index1)
+    #     self.model.train()
+    #     while loss.item() > loss_threshold and iteration < max_iter and (index1 < max_index) :
+    #         # print('iter', iteration)
+    #         x = x.requires_grad_(True)
+    #         # print('x shape', x.shape)
+    #         input['x'] = x
+    #         e_t,  att_first, att_second, att_third, self_first, self_second, self_third = self.model(input)
+            
+    #         bboxes = input['boxes_att']
+    #         object_positions = input['object_position'] 
+    #         loss2 = caculate_loss_LoCo(att_second,att_first,att_third, bboxes=bboxes,
+    #                     object_positions=object_positions, t = index1)*loss_scale
+    #         # loss = loss2
+    #         # loss.requires_grad_(True)
+    #         #print('LoCo loss', loss) 
+
+
+            
+    #         grad_cond = torch.autograd.grad(loss2.requires_grad_(True), [x])[0]
+    #         # grad_cond = x.grad 
+    #         x = x - grad_cond
+    #         x = x.detach()
+    #         iteration += 1
+    #         torch.cuda.empty_cache()
+    #     return x
+
+    def update_loss_LoCo(self, input,index1, index, ts,type_loss='self_accross'):
+       
         if index1 < 10:
-            loss_scale = 8
+            loss_scale = 3
             max_iter = 5
         elif index1 < 20:
-            loss_scale = 5
+            loss_scale = 2
             max_iter = 5
         else:
             loss_scale = 1
             max_iter = 1
         loss_threshold = 0.1
-        
+
         max_index = 30
         x = deepcopy(input["x"]) 
         iteration = 0
         loss = torch.tensor(10000)
         input["timesteps"] = ts
         
-        # print("optimize", index1)
-        self.model.train()
+        print("optimize", index1)
         while loss.item() > loss_threshold and iteration < max_iter and (index1 < max_index) :
-            # print('iter', iteration)
+            print('iter', iteration)
             x = x.requires_grad_(True)
-            # print('x shape', x.shape)
             input['x'] = x
             e_t,  att_first, att_second, att_third, self_first, self_second, self_third = self.model(input)
             
-            bboxes = input['boxes_att']
+            bboxes = input['boxes']
             object_positions = input['object_position'] 
             loss2 = caculate_loss_LoCo(att_second,att_first,att_third, bboxes=bboxes,
                         object_positions=object_positions, t = index1)*loss_scale
-            # loss = loss2
-            # loss.requires_grad_(True)
-            #print('LoCo loss', loss) 
-
-
-            
-            grad_cond = torch.autograd.grad(loss2.requires_grad_(True), [x])[0]
-            # grad_cond = x.grad 
-            x = x - grad_cond
+            loss = loss2
+            print('loss', loss) 
+            hh = torch.autograd.backward(loss)
+            grad_cond = x.grad
+            x = x - grad_cond 
             x = x.detach()
             iteration += 1
             torch.cuda.empty_cache()
         return x
-        
+    
     def update_loss_self_cross(self, input,index1, index, ts,type_loss='self_accross' ):
         if index1 < 10:
             loss_scale = 4