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

gligen/ldm/models/diffusion/plms.py

@@ -108,21 +108,10 @@ 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 )
-                    elif loss_type=='LoCo':
-
-                        #print('Utilizing LoCo!!')
-                        time_factor = noise_scheduler.sigmas[i] ** 2
-                        x = self.update_loss_LoCo( input,i, index, ts,  time_factor = time_factor)      
-
-                    elif loss_type=='LAC':
-                        #print('Utilizing LoCo!!')
-                        x = self.update_loss_LAC( input,i, index, ts )               
+
+                    time_factor = noise_scheduler.sigmas[i] ** 2
+                    x = self.update_loss_LoCo( input,i, index, ts,  time_factor = time_factor)      
+           
                     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
@@ -132,86 +121,7 @@ class PLMSSampler(object):
 
         return img
      
-    def update_loss_self_cross(self, input,index1, index, ts,type_loss='self_accross' ):
-        if index1 < 10:
-            loss_scale = 3
-            max_iter = 5
-        elif index1 < 20:
-            loss_scale = 2
-            max_iter = 3
-        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)
-            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'] 
-            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('AR loss:', loss, 'SAR:', loss1, 'CAR:', loss2)  
-            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_only_cross(self, input,index1, index, ts,type_loss='self_accross'):
-       
-        if index1 < 10:
-            loss_scale = 3
-            max_iter = 5
-        elif index1 < 20:
-            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)
-        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_att_fixed_cnt(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,  retain_graph=True)
-            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, time_factor, type_loss='self_accross'):
         
@@ -261,91 +171,8 @@ class PLMSSampler(object):
             torch.cuda.empty_cache()
         return x
 
-    def update_loss_LAC(self, input,index1, index, ts,type_loss='self_accross'):
-        
-        # loss_scale = 30
-        # max_iter = 5
-
-        if index1 < 10:
-            loss_scale = 6
-            max_iter = 5
-        elif index1 < 20:
-            loss_scale = 4
-            max_iter = 3
-        else:
-            loss_scale = 1
-            max_iter = 1
-        loss_threshold = 0.002
-
-        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_LAC(att_second,att_first,att_third, bboxes=bboxes,
-                        object_positions=object_positions, t = index1)*loss_scale
-            loss = loss2
-            print('LoCo loss', loss) 
-            hh = torch.autograd.backward(loss,  retain_graph=True)
-            grad_cond = x.grad
-            x = x - grad_cond 
-            x = x.detach()
-            iteration += 1
-            torch.cuda.empty_cache()
-        return x
-
 
 
-    def update_only_self(self, input,index1, index, ts,type_loss='self_accross' ):
-        if index1 < 10:
-            loss_scale = 4
-            max_iter = 5
-        elif index1 < 20:
-            loss_scale = 3
-            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)
-            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']
-            loss = caculate_loss_self_att(self_first, self_second, self_third, bboxes=bboxes,
-                                object_positions=object_positions, t = index1)*loss_scale 
-            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
-
     @torch.no_grad()
     def p_sample_plms(self, input, t, index, guidance_scale=1., uc=None, old_eps=None, t_next=None):
         x = deepcopy(input["x"])