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

gligen/ldm/models/diffusion/plms.py

@@ -3,7 +3,6 @@ import numpy as np
 from tqdm import tqdm
 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, caculate_loss_LoCo_V2
@@ -58,14 +57,14 @@ class PLMSSampler(object):
 
 
     # @torch.no_grad()
-    def sample(self, S, shape, input, uc=None, guidance_scale=1, mask=None, x0=None, loss_type=None):
+    def sample(self, S, shape, input, uc=None, guidance_scale=1, mask=None, x0=None, loss_type='LoCo'):
         self.make_schedule(ddim_num_steps=S)
         # import pdb; pdb.set_trace()
         return self.plms_sampling(shape, input, uc, guidance_scale, mask=mask, x0=x0, loss_type=loss_type)
 
 
     # @torch.no_grad()
-    def plms_sampling(self, shape, input, uc=None, guidance_scale=1, mask=None, x0=None, loss_type=None):
+    def plms_sampling(self, shape, input, uc=None, guidance_scale=1, mask=None, x0=None, loss_type='LoCo'):
 
         b = shape[0]
         
@@ -82,7 +81,6 @@ class PLMSSampler(object):
         if self.alpha_generator_func != None:
             alphas = self.alpha_generator_func(len(time_range))
 
-
         for i, step in enumerate(time_range):
 
             # set alpha and restore first conv layer 
@@ -104,7 +102,14 @@ class PLMSSampler(object):
             # three loss types
             if loss_type !=None and loss_type!='standard':
                 if input['object_position'] != []:
-                    x = self.update_loss_LoCo( input,i, index, ts,  time_factor = time_factor)              
+                    # 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':
+                    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
@@ -113,60 +118,67 @@ class PLMSSampler(object):
                 old_eps.pop(0)
 
         return img
-     
-    def update_loss_self_cross(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)
         if index1 < 10:
-            loss_scale = 3
+            loss_scale = 8
             max_iter = 5
         elif index1 < 20:
-            loss_scale = 2
-            max_iter = 3
+            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)
+        # 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']
             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
+            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) 
+
+
+            
+            hh = torch.autograd.backward(loss2,  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_only_cross(self, input,index1, index, ts,type_loss='self_accross'):
-       
+        
+    def update_loss_self_cross(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
@@ -174,68 +186,64 @@ 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)
-            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']
+            bboxes = input['boxes_att']
             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 = loss2
-            print('loss', loss) 
-            hh = torch.autograd.backward(loss,  retain_graph=True)
-            grad_cond = x.grad
-            x = x - grad_cond 
+                                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_LoCo(self, input,index1, index, ts, time_factor, type_loss='self_accross'):
-        
-        # loss_scale = 30
-        # max_iter = 5
-        #print('time_factor is: ', time_factor)
+    def update_loss_only_cross(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)
+        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']
             object_positions = input['object_position'] 
-            loss2 = caculate_loss_LoCo_V2(att_second,att_first,att_third, bboxes=bboxes,
+            loss2 = caculate_loss_att_fixed_cnt(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) 
-
-
-            
-            hh = torch.autograd.backward(loss2,  retain_graph=True)
+            loss = loss2
+            print('loss', loss) 
+            hh = torch.autograd.backward(loss)
             grad_cond = x.grad
             x = x - grad_cond 
             x = x.detach()
@@ -286,7 +294,7 @@ class PLMSSampler(object):
     def p_sample_plms(self, input, t, index, guidance_scale=1., uc=None, old_eps=None, t_next=None):
         x = deepcopy(input["x"]) 
         b = x.shape[0]
-
+        self.model.eval()
         def get_model_output(input):
             e_t, first, second, third,_,_,_ = self.model(input) 
             if uc is not None and guidance_scale != 1:
@@ -335,5 +343,3 @@ class PLMSSampler(object):
         x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
 
         return x_prev, pred_x0, e_t
-
-