remove save GPU memory

demo/model.py

@@ -175,11 +175,11 @@ class Model_all:
             im = im.float()
             im_edge = tensor2img(im)
         
-        # save gpu memory
-        self.base_model.model = self.base_model.model.cpu()
-        self.model_sketch = self.model_sketch.cuda()
-        self.base_model.first_stage_model = self.base_model.first_stage_model.cpu()
-        self.base_model.cond_stage_model = self.base_model.cond_stage_model.cuda()
+        # # save gpu memory
+        # self.base_model.model = self.base_model.model.cpu()
+        # self.model_sketch = self.model_sketch.cuda()
+        # self.base_model.first_stage_model = self.base_model.first_stage_model.cpu()
+        # self.base_model.cond_stage_model = self.base_model.cond_stage_model.cuda()
 
         # extract condition features
         c = self.base_model.get_learned_conditioning([prompt+', '+pos_prompt])
@@ -187,10 +187,10 @@ class Model_all:
         features_adapter = self.model_sketch(im.to(self.device))
         shape = [4, 64, 64]
         
-        # save gpu memory
-        self.model_sketch = self.model_sketch.cpu()
-        self.base_model.cond_stage_model = self.base_model.cond_stage_model.cpu()
-        self.base_model.model = self.base_model.model.cuda()
+        # # save gpu memory
+        # self.model_sketch = self.model_sketch.cpu()
+        # self.base_model.cond_stage_model = self.base_model.cond_stage_model.cpu()
+        # self.base_model.model = self.base_model.model.cuda()
 
         # sampling
         samples_ddim, _ = self.sampler.sample(S=50,
@@ -205,8 +205,8 @@ class Model_all:
                                         features_adapter1=features_adapter,
                                         mode = 'sketch',
                                         con_strength = con_strength)
-        # save gpu memory
-        self.base_model.first_stage_model = self.base_model.first_stage_model.cuda()
+        # # save gpu memory
+        # self.base_model.first_stage_model = self.base_model.first_stage_model.cuda()
 
         x_samples_ddim = self.base_model.decode_first_stage(samples_ddim)
         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
@@ -246,11 +246,11 @@ class Model_all:
         im = im>0.5
         im = im.float()
 
-        # save gpu memory
-        self.base_model.model = self.base_model.model.cpu()
-        self.model_sketch = self.model_sketch.cuda()
-        self.base_model.first_stage_model = self.base_model.first_stage_model.cpu()
-        self.base_model.cond_stage_model = self.base_model.cond_stage_model.cuda()
+        # # save gpu memory
+        # self.base_model.model = self.base_model.model.cpu()
+        # self.model_sketch = self.model_sketch.cuda()
+        # self.base_model.first_stage_model = self.base_model.first_stage_model.cpu()
+        # self.base_model.cond_stage_model = self.base_model.cond_stage_model.cuda()
         
         # extract condition features
         c = self.base_model.get_learned_conditioning([prompt+', '+pos_prompt])
@@ -258,10 +258,10 @@ class Model_all:
         features_adapter = self.model_sketch(im.to(self.device))
         shape = [4, 64, 64]
 
-        # save gpu memory
-        self.model_sketch = self.model_sketch.cpu()
-        self.base_model.cond_stage_model = self.base_model.cond_stage_model.cpu()
-        self.base_model.model = self.base_model.model.cuda()
+        # # save gpu memory
+        # self.model_sketch = self.model_sketch.cpu()
+        # self.base_model.cond_stage_model = self.base_model.cond_stage_model.cpu()
+        # self.base_model.model = self.base_model.model.cuda()
 
         # sampling
         samples_ddim, _ = self.sampler.sample(S=50,
@@ -277,8 +277,8 @@ class Model_all:
                                         mode = 'sketch',
                                         con_strength = con_strength)
         
-        # save gpu memory
-        self.base_model.first_stage_model = self.base_model.first_stage_model.cuda()
+        # # save gpu memory
+        # self.base_model.first_stage_model = self.base_model.first_stage_model.cuda()
 
         x_samples_ddim = self.base_model.decode_first_stage(samples_ddim)
         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
@@ -345,11 +345,11 @@ class Model_all:
                 thickness=2)
         im_pose = cv2.resize(im_pose,(512,512))
         
-        # save gpu memory
-        self.base_model.model = self.base_model.model.cpu()
-        self.model_pose = self.model_pose.cuda()
-        self.base_model.first_stage_model = self.base_model.first_stage_model.cpu()
-        self.base_model.cond_stage_model = self.base_model.cond_stage_model.cuda()
+        # # save gpu memory
+        # self.base_model.model = self.base_model.model.cpu()
+        # self.model_pose = self.model_pose.cuda()
+        # self.base_model.first_stage_model = self.base_model.first_stage_model.cpu()
+        # self.base_model.cond_stage_model = self.base_model.cond_stage_model.cuda()
 
         # extract condition features
         c = self.base_model.get_learned_conditioning([prompt+', '+pos_prompt])
@@ -358,10 +358,10 @@ class Model_all:
         pose = pose.unsqueeze(0)
         features_adapter = self.model_pose(pose.to(self.device))
 
-        # save gpu memory
-        self.model_pose = self.model_pose.cpu()
-        self.base_model.cond_stage_model = self.base_model.cond_stage_model.cpu()
-        self.base_model.model = self.base_model.model.cuda()
+        # # save gpu memory
+        # self.model_pose = self.model_pose.cpu()
+        # self.base_model.cond_stage_model = self.base_model.cond_stage_model.cpu()
+        # self.base_model.model = self.base_model.model.cuda()
 
         shape = [4, 64, 64]
 
@@ -379,8 +379,8 @@ class Model_all:
                                         mode = 'sketch',
                                         con_strength = con_strength)
 
-        # save gpu memory
-        self.base_model.first_stage_model = self.base_model.first_stage_model.cuda()
+        # # save gpu memory
+        # self.base_model.first_stage_model = self.base_model.first_stage_model.cuda()
 
         x_samples_ddim = self.base_model.decode_first_stage(samples_ddim)
         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)