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

gligen/ldm/models/diffusion/plms.py

@@ -108,7 +108,7 @@ class PLMSSampler(object):
                     #     x = self.update_only_self( input,i, index, ts )
                     # elif loss_type=='CAR':
                     #     x = self.update_loss_only_cross( input,i, index, ts )
-                    x = self.update_loss_only_cross( input,i, index, ts )
+                    x = self.update_loss_LoCo( input,i, index, ts )
                     input["x"] = x
             img, pred_x0, e_t = self.p_sample_plms(input, ts, index=index, uc=uc, guidance_scale=guidance_scale, old_eps=old_eps, t_next=ts_next)
             input["x"] = img
@@ -171,16 +171,16 @@ class PLMSSampler(object):
     def update_loss_LoCo(self, input,index1, index, ts,type_loss='self_accross'):
        
         if index1 < 10:
-            loss_scale = 3
-            max_iter = 5
+            loss_scale = 4
+            max_iter = 1
         elif index1 < 20:
-            loss_scale = 2
-            max_iter = 5
+            loss_scale = 3
+            max_iter = 1
         else:
             loss_scale = 1
             max_iter = 1
+        
         loss_threshold = 0.1
-
         max_index = 30
         x = deepcopy(input["x"]) 
         iteration = 0
@@ -188,24 +188,29 @@ class PLMSSampler(object):
         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)
+            # 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)
-            
-            bboxes = input['boxes']
+            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
-            print('loss', loss) 
-            hh = torch.autograd.backward(loss)
-            grad_cond = x.grad
-            x = x - grad_cond 
+            loss1 = caculate_loss_self_att(self_first, self_second, self_third, bboxes=bboxes,
+                                object_positions=object_positions, t = index1)*loss_scale 
+            loss2 = caculate_loss_att_fixed_cnt(att_second,att_first,att_third, bboxes=bboxes,
+                                object_positions=object_positions, t = index1)*loss_scale
+            loss = loss1 + loss2
+            print('loss', loss, loss1, loss2)  
+            # hh = torch.autograd.backward(loss, retain_graph=True)
+            grad_cond = torch.autograd.grad(loss.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_self_cross(self, input,index1, index, ts,type_loss='self_accross' ):