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

gligen/ldm/models/diffusion/plms.py

@@ -5,7 +5,7 @@ from functools import partial
 from copy import deepcopy
 from ldm.modules.diffusionmodules.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
 import math
-from ldm.models.diffusion.loss import  caculate_loss_att_fixed_cnt, caculate_loss_self_att
+from ldm.models.diffusion.loss import  caculate_loss_att_fixed_cnt, caculate_loss_self_att, caculate_loss_LoCo_V2
 class PLMSSampler(object):
     def __init__(self, diffusion, model, schedule="linear", alpha_generator_func=None, set_alpha_scale=None):
         super().__init__()
@@ -102,12 +102,13 @@ class PLMSSampler(object):
             # three loss types
             if loss_type !=None and loss_type!='standard':
                 if input['object_position'] != []:
-                    if loss_type=='SAR_CAR':
-                        x = self.update_loss_self_cross( input,i, index, ts )
-                    elif loss_type=='SAR':
-                        x = self.update_only_self( input,i, index, ts )
-                    elif loss_type=='CAR':
-                        x = self.update_loss_only_cross( input,i, index, ts )
+                    # if loss_type=='SAR_CAR':
+                    #     x = self.update_loss_self_cross( input,i, index, ts )
+                    # elif loss_type=='SAR':
+                    #     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_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
@@ -116,7 +117,56 @@ class PLMSSampler(object):
                 old_eps.pop(0)
 
         return img
-     
+
+
+    def update_loss_LoCo(self, input,index1, index, ts, time_factor, type_loss='self_accross'):
+        
+        # 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)
+        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']
+            object_positions = input['object_position'] 
+            loss2 = caculate_loss_LoCo_V2(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_self_cross(self, input,index1, index, ts,type_loss='self_accross' ):
         if index1 < 10:
             loss_scale = 4