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.")
-
-        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.")
+        # 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]
         samples = torch.cat((condition_img[0:1], samples), dim=0)
         samples = 255 * (samples * 0.5 + 0.5)
         samples = [
@@ -247,29 +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 = torch.cat((condition_img[0:1], samples), dim=0)
         samples = 255 * (samples * 0.5 + 0.5)
         samples = [