Update model.py

model.py

@@ -153,29 +153,29 @@ class Model:
         qzshape = [len(c_indices), 8, H // 16, W // 16]
         t1 = time.time()
         print(caption_embs.device)
-        # index_sample = generate(
-        #     self.gpt_model,
-        #     c_indices,
-        #     (H // 16) * (W // 16),
-        #     c_emb_masks,
-        #     condition=condition_img,
-        #     cfg_scale=cfg_scale,
-        #     temperature=temperature,
-        #     top_k=top_k,
-        #     top_p=top_p,
-        #     sample_logits=True,
-        #     control_strength=control_strength,
-        # )
-        # sampling_time = time.time() - t1
-        # print(f"Full sampling takes about {sampling_time:.2f} seconds.")
+        index_sample = generate(
+            self.gpt_model,
+            c_indices,
+            (H // 16) * (W // 16),
+            c_emb_masks,
+            condition=condition_img,
+            cfg_scale=cfg_scale,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            sample_logits=True,
+            control_strength=control_strength,
+        )
+        sampling_time = time.time() - t1
+        print(f"Full sampling takes about {sampling_time:.2f} seconds.")
 
-        # t2 = time.time()
-        # print(index_sample.shape)
-        # samples = self.vq_model.decode_code(
-        #     index_sample, qzshape)  # output value is between [-1, 1]
-        # decoder_time = time.time() - t2
-        # print(f"decoder takes about {decoder_time:.2f} seconds.")
-        samples = condition_img[0:1]
+        t2 = time.time()
+        print(index_sample.shape)
+        samples = self.vq_model.decode_code(
+            index_sample, qzshape)  # output value is between [-1, 1]
+        decoder_time = time.time() - t2
+        print(f"decoder takes about {decoder_time:.2f} seconds.")
+        # samples = condition_img[0:1]
         samples = torch.cat((condition_img[0:1], samples), dim=0)
         samples = 255 * (samples * 0.5 + 0.5)
         samples = [
@@ -247,31 +247,31 @@ class Model:
         c_emb_masks = new_emb_masks
         qzshape = [len(c_indices), 8, H // 16, W // 16]
         t1 = time.time()
-        # index_sample = generate(
-        #     self.gpt_model,
-        #     c_indices,
-        #     (H // 16) * (W // 16),
-        #     c_emb_masks,
-        #     condition=condition_img,
-        #     cfg_scale=cfg_scale,
-        #     temperature=temperature,
-        #     top_k=top_k,
-        #     top_p=top_p,
-        #     sample_logits=True,
-        #     control_strength=control_strength,
-        # )
-        # sampling_time = time.time() - t1
-        # print(f"Full sampling takes about {sampling_time:.2f} seconds.")
+        index_sample = generate(
+            self.gpt_model,
+            c_indices,
+            (H // 16) * (W // 16),
+            c_emb_masks,
+            condition=condition_img,
+            cfg_scale=cfg_scale,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            sample_logits=True,
+            control_strength=control_strength,
+        )
+        sampling_time = time.time() - t1
+        print(f"Full sampling takes about {sampling_time:.2f} seconds.")
 
-        # t2 = time.time()
-        # print(index_sample.shape)
-        # samples = self.vq_model.decode_code(index_sample, qzshape)
-        # decoder_time = time.time() - t2
-        # print(f"decoder takes about {decoder_time:.2f} seconds.")
-        # condition_img = condition_img.cpu()
-        # samples = samples.cpu()
+        t2 = time.time()
+        print(index_sample.shape)
+        samples = self.vq_model.decode_code(index_sample, qzshape)
+        decoder_time = time.time() - t2
+        print(f"decoder takes about {decoder_time:.2f} seconds.")
+        condition_img = condition_img.cpu()
+        samples = samples.cpu()
 
-        samples = condition_img[0:1]
+        # samples = condition_img[0:1]
         samples = torch.cat((condition_img[0:1], samples), dim=0)
         samples = 255 * (samples * 0.5 + 0.5)
         samples = [