Update model.py

model.py

@@ -122,8 +122,8 @@ class Model:
         elif preprocessor_name == 'No preprocess':
             condition_img = image
         print('get edge')
-        del self.preprocessor.model
-        torch.cuda.empty_cache()
+        # del self.preprocessor.model
+        # torch.cuda.empty_cache()
         condition_img = condition_img.resize((512,512))
         W, H = condition_img.size
 
@@ -131,10 +131,10 @@ class Model:
         self.load_gpt_weight('edge')
         self.gpt_model.to('cuda').to(torch.bfloat16)
         self.vq_model.to('cuda')
-        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(1,1,1,1)
+        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(3,1,1,1)
         condition_img = condition_img.to(self.device)
         condition_img = 2*(condition_img/255 - 0.5)
-        prompts = [prompt] * 1
+        prompts = [prompt] * 3
         caption_embs, emb_masks = self.t5_model.get_text_embeddings(prompts)
 
         print(f"processing left-padding...")
@@ -181,7 +181,7 @@ class Model:
         samples = [
             Image.fromarray(
                 sample.permute(1, 2, 0).cpu().detach().numpy().clip(
-                    0, 255).astype(np.uint8)) for sample in samples
+                    0, 255).astype(np.uint8).resize((origin_W, origin_H))) for sample in samples
         ]
         del condition_img
         torch.cuda.empty_cache()
@@ -225,10 +225,10 @@ class Model:
         self.load_gpt_weight('depth')
         self.gpt_model.to('cuda').to(torch.bfloat16)
         self.vq_model.to(self.device)
-        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(1,1,1,1)
+        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(3,1,1,1)
         condition_img = condition_img.to(self.device)
         condition_img = 2*(condition_img/255 - 0.5)
-        prompts = [prompt] * 1
+        prompts = [prompt] * 3
         caption_embs, emb_masks = self.t5_model.get_text_embeddings(prompts)
 
         print(f"processing left-padding...")
@@ -276,7 +276,7 @@ class Model:
         samples = 255 * (samples * 0.5 + 0.5)
         samples = [
             Image.fromarray(
-                sample.permute(1, 2, 0).cpu().detach().numpy().clip(0, 255).astype(np.uint8))
+                sample.permute(1, 2, 0).cpu().detach().numpy().clip(0, 255).astype(np.uint8).resize((origin_W, origin_H)))
             for sample in samples
         ]
         del condition_img